diff --git a/model/src/main.py b/model/src/main.py index 7069001..aff4a23 100644 --- a/model/src/main.py +++ b/model/src/main.py @@ -1,7 +1,12 @@ """Main script to be run by service.""" +import logging +from traceback import print_exc + +import torch 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.utils import setup_logging @@ -13,6 +18,29 @@ if __name__ == '__main__': minio_client = connect() # instantiate model model: VisualCommunicationModel = load_model(client=minio_client) - # get image + model.eval() + # setup dataset + dataset = VCDADataset( + data_name_list=[ + '02dbaf48d713e4e6d3a6b98fd2dc866e', + ], + do_augment=False, + ) # 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')