From 168ae8c44c80b44f895da1db8bf7cf85b165cff1 Mon Sep 17 00:00:00 2001
From: akhot2 <akhot2@illinois.edu>
Date: Fri, 14 Apr 2023 13:34:19 -0500
Subject: [PATCH] Update transfer/transfer.py
---
transfer/transfer.py | 15 +++++++++------
1 file changed, 9 insertions(+), 6 deletions(-)
diff --git a/transfer/transfer.py b/transfer/transfer.py
index a3f65da..d1ddd3c 100644
--- a/transfer/transfer.py
+++ b/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)
\ No newline at end of file
+ csvfile.write(csv)
--
GitLab