Files
visual_critical_discourse_a…/misc/prediction_upload_test.py
T
brian 312f525a46
Code Quality Pipeline / Check Code (pull_request) Failing after 2m25s
fixed argument for check_env function
2024-10-20 18:46:32 +00:00

47 lines
1.5 KiB
Python

from __future__ import annotations
from pathlib import Path
from dotenv import load_dotenv
from shared.data_store import connect_minio
from shared.database import connect_mongodb, upsert_prediction
from shared.database.classes import VisualCommunication
from shared.utils import check_env, setup_logging
from web_ui.src.main import NECESSARY_ENV_VAR_LIST
if __name__ == '__main__':
# load in env file
env_path = Path(__file__).parent.parent / 'server.env'
assert env_path.exists()
load_dotenv(env_path)
# ensure env vars set
check_env(NECESSARY_ENV_VAR_LIST)
# setup logging
setup_logging()
# connect to minIO
minio_client = connect_minio()
# connect to MongoDB
collection, db, client = connect_mongodb()
# 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, minio_client=minio_client)
for path in img_path_list
]
# generate random predictions
for vis_com in vis_com_list:
vis_com.generate_random_prediction()
# upload visual communication
for vis_com in vis_com_list:
if vis_com.prediction is None:
continue
upsert_prediction(
collection=collection,
vis_com_name=vis_com.name,
predictions=vis_com.prediction,
)