162 lines
3.8 KiB
Python
162 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from pymongo.collection import Collection
|
|
|
|
from .classes import NoDocumentFoundException
|
|
from .classes import VisualCommunication
|
|
from core.dto import ModelData
|
|
|
|
|
|
def count_documents(
|
|
collection: Collection,
|
|
has_annotation: bool = False,
|
|
) -> int:
|
|
"""
|
|
Get the total number of documents
|
|
in database that matches the filters.
|
|
"""
|
|
assert isinstance(collection, Collection)
|
|
assert isinstance(has_annotation, bool)
|
|
# build query
|
|
query = {}
|
|
if has_annotation:
|
|
query['annotation'] = {'$ne': None}
|
|
return collection.count_documents(filter=query)
|
|
|
|
|
|
def list_names(
|
|
collection: Collection,
|
|
has_annotation: bool = True,
|
|
) -> list[str]:
|
|
"""List the names of entries that match the filters."""
|
|
assert isinstance(collection, Collection)
|
|
assert isinstance(has_annotation, bool)
|
|
# build query
|
|
query = {}
|
|
if has_annotation:
|
|
query['annotation'] = {'$ne': None}
|
|
# execute query
|
|
res_list = collection.find(
|
|
filter=query,
|
|
projection={
|
|
'_id': False,
|
|
'name': True,
|
|
},
|
|
)
|
|
# extract information
|
|
name_list = [elem['name'] for elem in res_list]
|
|
return name_list
|
|
|
|
|
|
def total_documents(
|
|
collection: Collection,
|
|
) -> int:
|
|
"""Get total number of documents in database."""
|
|
return collection.count_documents(filter={})
|
|
|
|
|
|
def total_annotated(
|
|
collection: Collection,
|
|
) -> int:
|
|
"""Get total number of annotated documents in database."""
|
|
query = {
|
|
'annotation': {
|
|
'$ne': None,
|
|
},
|
|
}
|
|
return collection.count_documents(filter=query)
|
|
|
|
|
|
def get_visual_communication(
|
|
collection: Collection,
|
|
with_annotation: bool = False,
|
|
) -> VisualCommunication:
|
|
"""Get a random visual communication from the database."""
|
|
query = {}
|
|
if with_annotation:
|
|
query['annotation'] = {'$ne': None}
|
|
else:
|
|
query['annotation'] = {'$eq': None}
|
|
data = collection.aggregate([
|
|
{
|
|
'$match': query, # find using filters
|
|
},
|
|
{
|
|
'$sample': {
|
|
'size': 1, # get one random
|
|
},
|
|
},
|
|
])
|
|
data_list = list(data) # read data from cursor object
|
|
if len(data_list) == 0:
|
|
logging.error('failed getting visual communication')
|
|
raise NoDocumentFoundException()
|
|
logging.info('finished')
|
|
return VisualCommunication.model_validate(data_list[0])
|
|
|
|
|
|
def upsert_predictions(
|
|
collection: Collection,
|
|
vis_com_name: str,
|
|
predictions: ModelData,
|
|
) -> None:
|
|
"""Upsert prediction data in the database."""
|
|
query = {
|
|
'name': vis_com_name,
|
|
}
|
|
update = {
|
|
'$set': {
|
|
'prediction': predictions.model_dump(),
|
|
},
|
|
}
|
|
res = collection.update_one(
|
|
filter=query,
|
|
update=update,
|
|
upsert=True,
|
|
)
|
|
logging.debug('upserted document: %s', res)
|
|
logging.info('finished')
|
|
|
|
|
|
def upsert_annotations(
|
|
collection: Collection,
|
|
vis_com_name: str,
|
|
annotations: ModelData,
|
|
) -> None:
|
|
"""Upserts annotation data in the database."""
|
|
query = {
|
|
'name': vis_com_name,
|
|
}
|
|
update = {
|
|
'$set': {
|
|
'annotation': annotations.model_dump(),
|
|
},
|
|
}
|
|
res = collection.update_one(
|
|
filter=query,
|
|
update=update,
|
|
upsert=True,
|
|
)
|
|
logging.info('upserted document: %s', res)
|
|
logging.info('finished')
|
|
|
|
|
|
def upsert_visual_communication(
|
|
collection: Collection,
|
|
visual_communication_list: list[VisualCommunication],
|
|
) -> bool:
|
|
"""
|
|
Upsert VisualCommunication object in the database.
|
|
Returns bool stating success.
|
|
"""
|
|
response = collection.insert_many(
|
|
[
|
|
vis_com.model_dump()
|
|
for vis_com
|
|
in visual_communication_list
|
|
],
|
|
)
|
|
return response.acknowledged
|