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')