commit 548a50b31d07e475cba3e8404990bd10730c8a52
parent ab6da1cffacb38a40a0b439a5554bf84448bd2ad
Author: david cochran <about.trout@gmail.com>
Date: Wed, 21 Feb 2024 09:51:43 -0800
fixes #7 - turn BoatFinder into more general ObjectDetector
Diffstat:
6 files changed, 93 insertions(+), 91 deletions(-)
diff --git a/ppboatwatch/__init__.py b/ppboatwatch/__init__.py
@@ -1,4 +1,4 @@
from .stream_sampler import StreamSampler
-from .boat_finder import BoatFinder
+from .object_detector import ObjectDetector
-__all__ = ["StreamSampler", "BoatFinder"]
+__all__ = ["StreamSampler", "ObjectDetector"]
diff --git a/ppboatwatch/boat_finder.py b/ppboatwatch/boat_finder.py
@@ -1,40 +0,0 @@
-import logging
-import sys
-import torch
-
-from PIL import Image, ImageDraw
-from time import perf_counter
-from transformers import DetrImageProcessor, DetrForObjectDetection
-
-
-class BoatFinder:
- 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"
- self.processor = DetrImageProcessor.from_pretrained(
- "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({self.device}) initialized in {perf_counter() - t0} seconds"
- )
-
- def find(self, img):
- res = self.__match_results(img)
- for label, score, box in zip(res["labels"], res["scores"], res["boxes"]):
- label = self.model.config.id2label[label.item()]
- 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=self.thresh
- )[0]
diff --git a/ppboatwatch/find_boats.py b/ppboatwatch/find_boats.py
@@ -1,48 +0,0 @@
-import argparse
-import logging
-import os
-
-from PIL import Image, ImageDraw
-from .boat_finder import BoatFinder
-
-
-def main():
- 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:
- img_file = f"{base}/{f}"
- image = Image.open(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}")
diff --git a/ppboatwatch/find_objects.py b/ppboatwatch/find_objects.py
@@ -0,0 +1,50 @@
+import argparse
+import logging
+import os
+
+from PIL import Image, ImageDraw
+from .object_detector import ObjectDetector
+
+
+def main():
+ 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()
+
+ level = logging.INFO if args.verbose else logging.WARN
+ logging.basicConfig(level=level, format="%(asctime)s %(message)s")
+
+ 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)
+
+
+def detect_and_save(img_file, detector, outdir):
+ image = Image.open(img_file)
+ matches = list(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}")
diff --git a/ppboatwatch/object_detector.py b/ppboatwatch/object_detector.py
@@ -0,0 +1,40 @@
+import logging
+import sys
+import torch
+
+from PIL import Image, ImageDraw
+from time import perf_counter
+from transformers import DetrImageProcessor, DetrForObjectDetection
+
+
+class ObjectDetector:
+ 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"
+ self.processor = DetrImageProcessor.from_pretrained(
+ "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"ObjectDetector initialized for device {self.device}; took {perf_counter() - t0} seconds."
+ )
+
+ def find(self, img):
+ res = self.__match_results(img)
+ for label, score, box in zip(res["labels"], res["scores"], res["boxes"]):
+ label = self.model.config.id2label[label.item()]
+ 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=self.thresh
+ )[0]
diff --git a/pyproject.toml b/pyproject.toml
@@ -52,7 +52,7 @@ dependencies = [
]
[project.scripts]
-find-boats = "ppboatwatch.find_boats:main"
+find-objects = "ppboatwatch.find_objects:main"
sample-streams = "ppboatwatch.sample_streams:main"
[tool.setuptools]
