Update transfer/transfer.py

transfer/transfer.py

@@ -35,6 +35,11 @@ for file in os.listdir(dir):
         shutil.move(src_path, dest_path)
         save_paths.append(dest_path)
 
+#load model
+model_path = '../yolov5_model/runs/train/many_small_bugs_final/weights/best.pt'
+model = torch.hub.load('ultralytics/yolov5', 'custom', model_path)
+model.eval()
+
 # create output csv
 csv = 'Image Path, Location, Number of Beetles\n'
 for path in save_paths:
@@ -47,13 +52,11 @@ for path in save_paths:
     warped_rgb = crop(image)
 
     # run model
-    model_path = ''
-    model = torch.hub.load(model_path)
-    result = model(path)
-
+    result = model(warped_rgb)
+    num_beetles = results.pandas().xyxy[0].shape[0]
     # update csv
-    csv += f'{path},{trap_label},{result}\n'
+    csv += f'{path},{trap_label},{num_beetles}\n'
 
 csv_path = os.path.join(save_dir, 'output.csv')
 with open(csv_path, 'a') as csvfile:
-    csvfile.write(csv)
+    csvfile.write(csv)
\ No newline at end of file