model/src/main.py

@@ -1,5 +1,6 @@
 """Main script to be run by service."""
 
+import json
 import logging
 from traceback import print_exc
 
@@ -9,6 +10,7 @@ from tqdm import tqdm
 from utils import DEVICE, VCDADataset, load_model
 
 from shared.data_store import connect_minio
+from shared.dto import ModelData
 from shared.utils import setup_logging
 
 if __name__ == '__main__':
@@ -33,15 +35,14 @@ if __name__ == '__main__':
         for i in tqdm(range(len(dataset))):
             try:
                 # get image
-                image = dataset.__getitem__(i)
+                image = dataset[i]
                 image = torch.unsqueeze(image, 0)  # add artificial batch dimension
                 image = image.to(DEVICE)
                 # make prediction
-                pred = model(image)
-                pred = pred.cpu().numpy()
+                pred: ModelData = model(image)
             except Exception:
                 print_exc()
                 continue
             else:
-                logging.info(pred)
+                print(json.dumps(pred.model_dump(), indent=4))
     logging.debug('finished')