diff --git a/__pycache__/main.cpython-39.pyc b/__pycache__/main.cpython-39.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..712512dcc36dce0b1150fc11d1519c691587d4be
Binary files /dev/null and b/__pycache__/main.cpython-39.pyc differ
diff --git a/crop/__init__.py b/crop/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/crop/__pycache__/__init__.cpython-39.pyc b/crop/__pycache__/__init__.cpython-39.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..23734692590ae5493f852f2fc189965818cfe180
Binary files /dev/null and b/crop/__pycache__/__init__.cpython-39.pyc differ
diff --git a/crop/__pycache__/autocrop.cpython-39.pyc b/crop/__pycache__/autocrop.cpython-39.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..301268a0174bf405d153e3e6e3e6cf1a6439a015
Binary files /dev/null and b/crop/__pycache__/autocrop.cpython-39.pyc differ
diff --git a/crop/autocrop.py b/crop/autocrop.py
index dd1f6acf47d7d4f24f83a83c91ff8abaaf0cfbe1..90995d6cc08e0f0605f01ffe92c473cdcc1c109a 100644
--- a/crop/autocrop.py
+++ b/crop/autocrop.py
@@ -1,8 +1,6 @@
# 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
diff --git a/main.py b/main.py
new file mode 100644
index 0000000000000000000000000000000000000000..41bd01d1b15111add8b4a1a75f5a63a0151c83bd
--- /dev/null
+++ b/main.py
@@ -0,0 +1 @@
+import transfer.transfer
\ No newline at end of file
diff --git a/transfer/__init__.py b/transfer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/transfer/__pycache__/__init__.cpython-39.pyc b/transfer/__pycache__/__init__.cpython-39.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bf43265181abad488699d75cf62cafc472a7f20c
Binary files /dev/null and b/transfer/__pycache__/__init__.cpython-39.pyc differ
diff --git a/transfer/__pycache__/transfer.cpython-39.pyc b/transfer/__pycache__/transfer.cpython-39.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..278ae665211d8fccca80e0299449b1250c3dd623
Binary files /dev/null and b/transfer/__pycache__/transfer.cpython-39.pyc differ
diff --git a/transfer/output.csv b/transfer/output.csv
new file mode 100644
index 0000000000000000000000000000000000000000..a09207ae47b8ba26e558041bb2199c92158f7951
--- /dev/null
+++ b/transfer/output.csv
@@ -0,0 +1,4 @@
+Image Path, Location, Number of Beetles
+test
+test
+test3
diff --git a/transfer/transfer.ipynb b/transfer/transfer.ipynb
index 0039c8bb3803908f2515709bb08ff64ac638929f..8f2a0df1a44f8a43ea5e28e3f10431a601ca7900 100644
--- a/transfer/transfer.ipynb
+++ b/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",
diff --git a/transfer/transfer.py b/transfer/transfer.py
index d1ddd3c93224e571a66b377bb4f2b5c23a7060a4..6ed2a482bb4afac280af8a1f6b4fd7aa8f8e5182 100644
--- a/transfer/transfer.py
+++ b/transfer/transfer.py
@@ -1,11 +1,13 @@
-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