PyCS/models/moth_scanner/scanner/detector.py

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