Files
visual_critical_discourse_a…/tests/prediction_upload_test.py
T
2024-02-25 14:26:54 +01:00

52 lines
1.4 KiB
Python

from __future__ import annotations
import logging
import os
from pathlib import Path
from database import connect
from database import upsert_predictions
from database import VisualCommunication
from dotenv import load_dotenv
if __name__ == '__main__':
# setup logging
fmt = (
'%(asctime)s | '
'%(levelname)s | '
'%(filename)s | '
'%(funcName)s | '
'%(message)s'
)
datefmt = '%Y-%m-%d %H:%M:%S'
logging.basicConfig(format=fmt, datefmt=datefmt, level=logging.INFO)
# get list of image paths
test_dir = Path(__file__).parent
img_dir = test_dir / 'imgs'
img_path_list = [path for path in img_dir.glob('*.jpeg') if path.is_file()]
# instantiate data object
vis_com_list = [
VisualCommunication.from_file(path)
for path
in img_path_list
]
# generate random predictions
for vis_com in vis_com_list:
vis_com.generate_random_prediction()
# prepare env vars
env_path = test_dir.parent / 'local.env'
assert env_path.exists()
load_dotenv(env_path)
os.environ['MONGO_HOST'] = 'localhost'
# connect to database
collection, db, client = connect()
# upload visual communication
for vis_com in vis_com_list:
if vis_com.prediction is None:
continue
upsert_predictions(
collection=collection,
vis_com_name=vis_com.name,
predictions=vis_com.prediction,
)