test run prediction on image

This commit is contained in:
Brian Bjarke Jensen
2024-07-27 21:41:29 +02:00
parent 3ad6c161b1
commit 6c354b23bd
+30 -2
View File
@@ -1,7 +1,12 @@
"""Main script to be run by service.""" """Main script to be run by service."""
import logging
from traceback import print_exc
import torch
from models import VisualCommunicationModel from models import VisualCommunicationModel
from utils import load_model from tqdm import tqdm
from utils import DEVICE, VCDADataset, load_model
from shared.data_store import connect from shared.data_store import connect
from shared.utils import setup_logging from shared.utils import setup_logging
@@ -13,6 +18,29 @@ if __name__ == '__main__':
minio_client = connect() minio_client = connect()
# instantiate model # instantiate model
model: VisualCommunicationModel = load_model(client=minio_client) model: VisualCommunicationModel = load_model(client=minio_client)
# get image model.eval()
# setup dataset
dataset = VCDADataset(
data_name_list=[
'02dbaf48d713e4e6d3a6b98fd2dc866e',
],
do_augment=False,
)
# make prediction # make prediction
with torch.no_grad():
for i in tqdm(range(len(dataset))):
try:
# get image
image = dataset.__getitem__(i)
image = torch.unsqueeze(image, 0) # add artificial batch dimension
image = image.to(DEVICE)
# make prediction
pred = model(image)
pred = pred.cpu().numpy()
except Exception:
print_exc()
continue
else:
logging.info(pred)
logging.debug('finished')