121 lines
2.8 KiB
Python
121 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from pymongo.collection import Collection
|
|
|
|
from .classes import ModelOutputs
|
|
from .classes import NoDocumentFoundException
|
|
from .classes import VisualCommunication
|
|
|
|
|
|
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: ModelOutputs,
|
|
) -> 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: ModelOutputs,
|
|
) -> 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
|