ppbbww

commit 85677236d31ac137928d8b7989e68bda6672bc03
parent 857c12f33d5cfc3aff16111f8d1ce00d8eb09c3b
Author: david cochran <about.trout@gmail.com>
Date:   Sat, 30 Mar 2024 19:16:07 -0700

archive matches rather than writing copies

Diffstat:
M
ppboatwatch/find_objects.py
 | 
29
+++++++++++------------------

1 file changed, 11 insertions(+), 18 deletions(-)
diff --git a/ppboatwatch/find_objects.py b/ppboatwatch/find_objects.py
@@ -3,7 +3,10 @@ import logging
 import os
 
 from PIL import Image, ImageDraw
+
+from .archive import Archive
 from .object_detector import ObjectDetector
+from .ppbw import filter_matches
 
 
 def main():
@@ -11,40 +14,30 @@ def main():
     parser.add_argument("filepath")
     parser.add_argument("-t", "--thresh", type=float, default=0.9)
     parser.add_argument("-l", "--labels", type=str, nargs="+", action="extend")
-    parser.add_argument(
-        "-o",
-        "--outdir",
-        type=str,
-        help="Save annotated matches to given folder, if set.",
-    )
+    parser.add_argument("-f", "--db-file", default="./archive.db", type=str)
     parser.add_argument("-v", "--verbose", action="store_true")
     args = parser.parse_args()
 
     level = logging.INFO if args.verbose else logging.WARN
     logging.basicConfig(level=level, format="%(asctime)s %(message)s")
 
+    archive = Archive(args.db_file)
     detector = ObjectDetector(thresh=args.thresh, labels=args.labels)
     for base, _, fs in os.walk(args.filepath):
         for f in fs:
-            detect_and_save(f"{base}/{f}", detector, args.outdir)
+            detect_and_save(f"{base}/{f}", detector, archive)
 
 
-def detect_and_save(img_file, detector, outdir):
+def detect_and_save(img_file, detector, archive):
     image = Image.open(img_file)
-    matches = list(detector.find(image))
+    matches = list(filter_matches(detector.find(image)))
     if len(matches) == 0:
         logging.info(f"No matches in {img_file}.")
         return
 
     logging.warn(f"Found {len(matches)} in {img_file}.")
-    draw = ImageDraw.Draw(image)
     for label, score, box in matches:
         logging.warn(f" >> {label} ({score})")
-        if outdir:
-            box = list(box)
-            draw.text((box[0], box[1] - 15), label, (166, 226, 46))
-            draw.rectangle(box, outline="#a6e22e", width=2)
-
-    if outdir:
-        out_file = img_file.split("/")[-1].removesuffix(".jpg") + "_matches.jpg"
-        image.save(f"{outdir}/{out_file}")
+        box = list(box)
+        ts = img_file.split("/")[-1].split("-")[0]
+        archive.add_match(ts=ts, filename=img_file, label=label, score=score, box=box)