commit 6352b48db4f62f22b16a22482e7e2912bb2b45c0
parent a2d391d38bce9ffa2f498a63aafd470f6b3ff45f
Author: david cochran <about.trout@gmail.com>
Date: Wed, 14 Feb 2024 13:57:14 -0800
fixes #3 - draw bounding boxes with PIL.ImageDraw
Diffstat:
2 files changed, 20 insertions(+), 19 deletions(-)
diff --git a/README.md b/README.md
@@ -2,7 +2,7 @@
Locate large boats as they pass by Pillar Point.
-
+
Install the requirements.
@@ -32,18 +32,11 @@ $ ./streamsampler.py
Find boats in images with [`BoatFinder`](./boatfinder.py#L12).
```
-$ ./boatfinder.py $(find ./data/mavericks/02102024/10/*.jpg | head -n10)
-2024-02-10 10:40:43,803 BoatFinder initializing ...
-2024-02-10 10:40:44,473 BoatFinder initialized! took 0.6702592820074642 seconds
-2024-02-10 10:40:44,477 Searching file ./data/mavericks/02102024/10/1707589172-thumb-0001.jpg...
-2024-02-10 10:40:46,057 Searching file ./data/mavericks/02102024/10/1707589172-thumb-0002.jpg...
-2024-02-10 10:40:47,474 Searching file ./data/mavericks/02102024/10/1707589172-thumb-0003.jpg...
-2024-02-10 10:40:48,865 Searching file ./data/mavericks/02102024/10/1707589172-thumb-0004.jpg...
-2024-02-10 10:40:50,264 Searching file ./data/mavericks/02102024/10/1707589172-thumb-0005.jpg...
-2024-02-10 10:40:51,669 Searching file ./data/mavericks/02102024/10/1707589250-thumb-0001.jpg...
-2024-02-10 10:40:53,074 >> label=boat score=0.677 box=[431.265, 673.42, 694.994, 719.792]
-2024-02-10 10:40:53,076 Searching file ./data/mavericks/02102024/10/1707589250-thumb-0002.jpg...
-2024-02-10 10:40:54,760 Searching file ./data/mavericks/02102024/10/1707589250-thumb-0003.jpg...
-2024-02-10 10:40:56,160 Searching file ./data/mavericks/02102024/10/1707589250-thumb-0004.jpg...
-2024-02-10 10:40:57,594 Searching file ./data/mavericks/02102024/10/1707589250-thumb-0005.jpg...
+$ ./boatfinder.py ./data/mavericksov/02142024/11/1707939852-thumb-0003.jpg
+2024-02-14 13:56:06,022 BoatFinder initializing ...
+2024-02-14 13:56:08,380 BoatFinder initialized (for cuda:0); took 2.3576930790004553 seconds
+2024-02-14 13:56:08,384 Searching file ./data/mavericksov/02142024/11/1707939852-thumb-0003.jpg...
+2024-02-14 13:56:09,939 Finished search in 1.5548269860009896 seconds
+2024-02-14 13:56:09,939 >> label=boat score=0.994 box=[517, 81, 910, 147]
+2024-02-14 13:56:09,944 Outlined matches and saved to file ./data/mavericksov/02142024/11/1707939852-thumb-0003_boats.jpg
```
diff --git a/boatfinder.py b/boatfinder.py
@@ -4,7 +4,7 @@ import logging
import sys
import torch
-from PIL import Image
+from PIL import Image, ImageDraw
from time import perf_counter
from transformers import DetrImageProcessor, DetrForObjectDetection
@@ -30,7 +30,7 @@ class BoatFinder:
label = self.model.config.id2label[label.item()]
if label == "boat":
score = round(score.item(), 3)
- box = [round(i, 3) for i in box.tolist()]
+ box = [int(i) for i in box.tolist()]
yield (score, label, box)
logging.info(f"Finished search in {perf_counter() - t0} seconds")
@@ -54,5 +54,13 @@ if __name__ == "__main__":
for img_file in sys.argv[1:]:
image = Image.open(img_file)
logging.info(f"Searching file {img_file}...")
- for score, label, box in bf.find(image):
- logging.info(f">> label={label}\t score={score}\t box={box}\t")
+ res = list(bf.find(image))
+ if len(res) > 0:
+ draw = ImageDraw.Draw(image)
+ for score, label, box in res:
+ logging.info(f">> label={label}\t score={score}\t box={box}\t")
+ draw.rectangle(box, outline="#a6e22e")
+ out_file = img_file.removesuffix(".jpg") + "_boats.jpg"
+ image.save(out_file)
+ logging.info(f"Outlined matches and saved to file {out_file}")
+
