| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768 |
- 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.rescale(
- min_size=config.preprocess.min_size,
- min_scale=config.preprocess.scale)
- self.img_proc.find_border()
- self.img_proc.preprocess(
- equalize=False, sigma=config.preprocess.sigma)
- self.img_proc.binarize(
- type=BinarizerType.gauss_local,
- use_masked=True,
- use_cv2=True,
- sigma=5.0,
- 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()
- _, 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)
- bboxes, labels, scores = self.bbox_proc(_im)
- return bboxes
|
