find_objects.py (1497B)
1 import argparse 2 import logging 3 import os 4 5 from PIL import Image, ImageDraw 6 7 from .archive import Archive 8 from .object_detector import ObjectDetector 9 from .ppbw import filter_matches 10 11 12 def main(): 13 parser = argparse.ArgumentParser() 14 parser.add_argument("filepath") 15 parser.add_argument("-t", "--thresh", type=float, default=0.9) 16 parser.add_argument("-l", "--labels", type=str, nargs="+", action="extend") 17 parser.add_argument("-f", "--db-file", default="./archive.db", type=str) 18 parser.add_argument("-v", "--verbose", action="store_true") 19 args = parser.parse_args() 20 21 level = logging.INFO if args.verbose else logging.WARN 22 logging.basicConfig(level=level, format="%(asctime)s %(message)s") 23 24 archive = Archive(args.db_file) 25 detector = ObjectDetector(thresh=args.thresh, labels=args.labels) 26 for base, _, fs in os.walk(args.filepath): 27 for f in fs: 28 detect_and_save(f"{base}/{f}", detector, archive) 29 30 31 def detect_and_save(img_file, detector, archive): 32 image = Image.open(img_file) 33 matches = list(filter_matches(detector.find(image))) 34 if len(matches) == 0: 35 logging.info(f"No matches in {img_file}.") 36 return 37 38 logging.warn(f"Found {len(matches)} in {img_file}.") 39 for label, score, box in matches: 40 logging.warn(f" >> {label} ({score})") 41 box = list(box) 42 ts = img_file.split("/")[-1].split("-")[0] 43 archive.add_match(ts=ts, filename=img_file, label=label, score=score, box=box)