import cv2 import numpy as np import typing as T from munch import munchify from collections import namedtuple from blob_detector import utils from blob_detector.core.bbox import BBox from blob_detector.core.bbox_proc import Splitter from blob_detector.core.binarizers import BinarizerType from blob_detector.core.pipeline import Pipeline class Detector(object): def __init__(self, configuration: T.Dict[str, T.Dict]) -> None: super().__init__() config = munchify(configuration) self.img_proc = Pipeline() self.img_proc.find_border() if config.preprocess.min_size: self.img_proc.rescale( min_size=config.preprocess.min_size, min_scale=config.preprocess.scale) self.img_proc.preprocess( equalize=False, sigma=config.preprocess.sigma) self.img_proc.binarize( type=BinarizerType.gauss_local, use_masked=True, use_cv2=True, window_size=config.threshold.window_size, offset=2.0, ) self.img_proc.remove_border() self.img_proc.open_close( kernel_size=config.postprocess.kernel_size, iterations=config.postprocess.dilate_iterations) self.bbox_proc = Pipeline() self.bbox_proc.detect(use_masked=True) _, splitter = self.bbox_proc.split_bboxes( preproc=Pipeline(), detector=Pipeline()) _, bbox_filter = self.bbox_proc.bbox_filter( score_threshold=0.5, nms_threshold=0.3, enlarge=0.01, ) # _, scorer = self.bbox_proc.score() self.img_proc.requires_input(splitter.set_image) self.img_proc.requires_input(bbox_filter.set_image) # self.img_proc.requires_input(scorer.set_image) def __call__(self, im: np.ndarray) -> T.List[BBox]: _im = self.img_proc(im) det_result = self.bbox_proc(_im) return det_result.bboxes