ppbbww

commit ab6da1cffacb38a40a0b439a5554bf84448bd2ad
parent a125eb996cc8a411492fc8c7d2ce96dafa37acef
Author: david cochran <about.trout@gmail.com>
Date:   Mon, 19 Feb 2024 20:12:57 -0800

add command line args to configure searches

Diffstat:
M
ppboatwatch/boat_finder.py
 | 
30
+++++++++++++++---------------
M
ppboatwatch/find_boats.py
 | 
52
+++++++++++++++++++++++++++++++++++++---------------

2 files changed, 52 insertions(+), 30 deletions(-)
diff --git a/ppboatwatch/boat_finder.py b/ppboatwatch/boat_finder.py
@@ -8,33 +8,33 @@ from transformers import DetrImageProcessor, DetrForObjectDetection
 
 
 class BoatFinder:
-    def __init__(self):
-        logging.info(f"BoatFinder initializing ...")
+    def __init__(self, thresh=0.95, labels=None):
+        self.thresh = thresh
+        self.filter_labels = labels
+
         t0 = perf_counter()
         self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        # https://huggingface.co/facebook/detr-resnet-50
         self.processor = DetrImageProcessor.from_pretrained(
-            "facebook/detr-resnet-50", revision="no_timm")
+            "facebook/detr-resnet-50", revision="no_timm"
+        )
         self.model = DetrForObjectDetection.from_pretrained(
-            "facebook/detr-resnet-50", revision="no_timm").to(self.device)
-        logging.info(f"BoatFinder initialized (for {self.device}); took {perf_counter() - t0} seconds")
+            "facebook/detr-resnet-50", revision="no_timm"
+        ).to(self.device)
+        logging.info(
+            f"BoatFinder({self.device}) initialized in {perf_counter() - t0} seconds"
+        )
 
     def find(self, img):
-        t0 = perf_counter()
         res = self.__match_results(img)
-        for score, label, box in zip(res["scores"], res["labels"], res["boxes"]):
+        for label, score, box in zip(res["labels"], res["scores"], res["boxes"]):
             label = self.model.config.id2label[label.item()]
-            if label == "boat":
-                score = round(score.item(), 3)
-                box = [int(i) for i in box.tolist()]
-                yield (score, label, box)
-        logging.info(f"Finished search in {perf_counter() - t0} seconds")
+            if (not self.filter_labels) or (label in self.filter_labels):
+                yield (label, round(score.item(), 2), map(int, box.tolist()))
 
     def __match_results(self, image):
         inputs = self.processor(images=image, return_tensors="pt").to(self.device)
         outputs = self.model(**inputs)
         target_sizes = torch.tensor([image.size[::-1]]).to(self.device)
         return self.processor.post_process_object_detection(
-            outputs, target_sizes=target_sizes, threshold=0.5
+            outputs, target_sizes=target_sizes, threshold=self.thresh
         )[0]
-
diff --git a/ppboatwatch/find_boats.py b/ppboatwatch/find_boats.py
@@ -1,3 +1,4 @@
+import argparse
 import logging
 import os
 
@@ -6,21 +7,42 @@ from .boat_finder import BoatFinder
 
 
 def main():
-    logging.basicConfig(level=logging.WARN, format="%(asctime)s %(message)s")
-    bf = BoatFinder()
-    find_dir = "./data"  # TODO: argparse.
-    for base, _, fs in os.walk(find_dir):
+    parser = argparse.ArgumentParser()
+    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("-v", "--verbose", action="store_true")
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.INFO if args.verbose else logging.WARN,
+        format="%(asctime)s %(message)s",
+    )
+
+    bf = BoatFinder(thresh=args.thresh, labels=args.labels)
+    for base, _, fs in os.walk(args.filepath):
         for f in fs:
-            # TODO: Ignore _boats.jpg files!
             img_file = f"{base}/{f}"
             image = Image.open(img_file)
-            res = list(bf.find(image))
-            if len(res) > 0:
-                draw = ImageDraw.Draw(image)
-                for score, label, box in res: 
-                    draw.rectangle(box, outline="#ff0000")
-                out_file = img_file.removesuffix(".jpg") + "_boats.jpg"
-                image.save(out_file)
-                logging.warning(f"Found {len(res)} matches in {img_file}")
-            else:
-                logging.warning(f"No matches in {img_file}")
+            matches = list(bf.find(image))
+            if len(matches) == 0:
+                logging.info(f"No matches in {img_file}.")
+                continue
+
+            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 args.outdir:
+                    # TODO: Add label text. Align top/left.
+                    draw.rectangle(list(box), outline="#a6e22e", width=2)
+
+            if args.outdir:
+                out_file = f.removesuffix(".jpg") + "_matches.jpg"
+                image.save(f"{args.outdir}/{out_file}")