Update eval.py

eval.py

@@ -17,35 +17,35 @@ if __name__ == '__main__':
     models.append(model)
 
     images = []
-    transform = T.Compose([
-        T.ToTensor(),
-        T.Resize((480, 640))
-    ])
-
     base_dir = 'test_images'
     for f in os.listdir(base_dir):
         if not f.endswith('.bmp'):
             continue
-        image = Image.open(f'{base_dir}/{f}').convert('RGB')
-        image = transform(image).unsqueeze(0)
+        image = Image.open(f'{base_dir}/{f}').convert('RGB').resize((640, 480))
         images.append(image)
 
     # Convert to OpenVINO
+    to_tensor = T.ToTensor()
     print('Converting to OpenVINO...')
     core = Core()
-    model = convert_model(model, example_input=images[0])
+    model = convert_model(model, example_input=to_tensor(images[0]).unsqueeze(0))
     model = core.compile_model(model, 'AUTO')
     models.append(model)
     print('Converted to OpenVINO.')
 
     for model in models:
         print(type(model))
+        count = 0
         for image in images:
             start_time = timeit.default_timer()
             if isinstance(model, CompiledModel):
-                preds = model(image)
+                preds = model(to_tensor(image).unsqueeze(0))
+                # TODO: Decode model output
+                # refer super_gradients.training.pipelines.pipelines -> DetectionPipeline._decode_model_output
             elif isinstance(model, Sg_models.SgModule):
                 preds = model.predict(image)
+                count += len(preds[0].prediction)
             elif isinstance(model, nn.Module):
                 preds = model(image)
-            print(f'Time: {(timeit.default_timer() - start_time) * 100:.3f}ms')
+            print(f'Time: {(timeit.default_timer() - start_time) * 100:.3f}ms')
+        print(f'Count: {count}')