main.py

crop/autocrop.py

 # https://stackoverflow.com/questions/60941012/how-do-i-find-corners-of-a-paper-when-there-are-printed-corners-lines-on-paper-i
 
 from PIL import Image
-from pillow_heif import register_heif_opener
-register_heif_opener()
 import numpy as np
 import cv2
 import sys

main.py

+import transfer.transfer
\ No newline at end of file

transfer/output.csv

+Image Path, Location, Number of Beetles
+test
+test
+test3

transfer/transfer.ipynb

@@ -2,23 +2,11 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 11,
    "metadata": {},
-   "outputs": [
-    {
-     "ename": "ModuleNotFoundError",
-     "evalue": "No module named 'torch'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[5], line 11\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mcrop\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mautocrop\u001b[39;00m \u001b[39mimport\u001b[39;00m crop\n\u001b[1;32m     10\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mcv2\u001b[39;00m\n\u001b[0;32m---> 11\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtorch\u001b[39;00m\n",
-      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'torch'"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "import os, sys\n",
+    "import os, sys, io\n",
     "dir2 = os.path.abspath('')\n",
     "dir1 = os.path.dirname(dir2)\n",
     "if not dir1 in sys.path: sys.path.append(dir1)\n",
@@ -126,10 +114,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "csv_title = 'Image Path, Location, Number of Beetles\\n'\n",
+    "csv = 'test3\\n'\n",
+    "csv_path = os.path.join('.', 'output.csv')\n",
+    "with open(csv_path, 'a+') as csvfile:\n",
+    "    csvfile.seek(0)\n",
+    "    if csvfile.readline().strip() == csv_title.strip():\n",
+    "        csvfile.seek(0, io.SEEK_END)\n",
+    "    else:\n",
+    "        csv = csv_title + csv\n",
+    "    csvfile.write(csv)"
+   ]
   },
   {
    "cell_type": "code",

transfer/transfer.py

-import os
+import os, io
 import shutil
 from tkinter.filedialog import askdirectory
 from tkinter.messagebox import askyesno
-from ..crop.autocrop import crop
+from crop.autocrop import crop
 import cv2
+from PIL import Image
 import torch
 
+
 def readQR(image : cv2.Mat):
     qrCodeDetector = cv2.QRCodeDetector()
     decoded_text, points, _ = qrCodeDetector.detectAndDecode(image)
@@ -13,7 +15,6 @@ def readQR(image : cv2.Mat):
         return ''
     return decoded_text
 
-
 # find sd path
 sd_path = askdirectory(message='Select Path To SD')
 assert sd_path
@@ -28,7 +29,7 @@ save_paths = []
 dir = os.fsencode(sd_path)
 for file in os.listdir(dir):
     filename = os.fsdecode(file)
-    if filename.endswith(".png"):
+    if filename.endswith(".png") or filename.endswith(".jpg"):
         # save original image to computer
         src_path = os.path.join(sd_path, filename)
         dest_path = os.path.join(save_dir, filename)
@@ -36,27 +37,36 @@ for file in os.listdir(dir):
         save_paths.append(dest_path)
 
 #load model
-model_path = '../yolov5_model/runs/train/many_small_bugs_final/weights/best.pt'
+model_path = 'yolov5_model/runs/train/20beetle_40-non_20dirt_bkg_overlap/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'
+csv_title = 'Image Path, Location, Number of Beetles, Confidence\n'
+csv = ''
 for path in save_paths:
     image = cv2.imread(path)
     
     # read QR
     trap_label = readQR(image)
 
-    # crop
+    # crop + save cropped
     warped_rgb = crop(image)
+    Image.fromarray(warped_rgb).save(path+"cropped.jpg")
 
     # run model
     result = model(warped_rgb)
-    num_beetles = results.pandas().xyxy[0].shape[0]
+    num_beetles = result.pandas().xyxy[0].shape[0]
+    confidence = result.pandas().xyxy[0]['confidence'].sum()
     # update csv
-    csv += f'{path},{trap_label},{num_beetles}\n'
+    csv += f'{path},{trap_label},{num_beetles},{confidence}\n'
 
+# save to csv
 csv_path = os.path.join(save_dir, 'output.csv')
-with open(csv_path, 'a') as csvfile:
-    csvfile.write(csv)
+with open(csv_path, 'a+') as csvfile:
+    csvfile.seek(0)
+    if csvfile.readline().strip() == csv_title.strip():
+        csvfile.seek(0, io.SEEK_END)
+    else:
+        csv = csv_title + csv
+    csvfile.write(csv)
\ No newline at end of file