diff --git a/transfer/transfer.py b/transfer/transfer.py index a3f65da3f9d5ba41abc90b26578cb0dd2a3be2b8..d1ddd3c93224e571a66b377bb4f2b5c23a7060a4 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)