91 lines
2.2 KiB
Python
91 lines
2.2 KiB
Python
from pymongo.collection import Collection
|
|
import logging
|
|
|
|
from .classes import (
|
|
VisualCommunication,
|
|
NoDocumentFoundException,
|
|
ModelOutputs
|
|
)
|
|
|
|
|
|
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"] = None
|
|
data = collection.aggregate([
|
|
{ "$match": query }, # find using filters
|
|
{ "$sample": { "size": 1 } } # get one random
|
|
])
|
|
data = list(data) # read data from cursor object
|
|
if len(data) == 0:
|
|
logging.error("failed getting visual communication")
|
|
raise NoDocumentFoundException()
|
|
data = data[0]
|
|
logging.info("finished")
|
|
return VisualCommunication.model_validate(data)
|
|
|
|
|
|
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")
|