PyCS/models/moth_scanner/scanner/__init__.py

import cv2
import numpy as np

from json import dump, load

from pycs.interfaces.MediaFile import MediaFile
from pycs.interfaces.MediaStorage import MediaStorage
from pycs.interfaces.Pipeline import Pipeline as Interface

from .detector import Detector
from .classifier import Classifier

class Scanner(Interface):
    def __init__(self, root_folder: str, configuration: dict):
        super().__init__(root_folder, configuration)
        self.detector = Detector(configuration["detector"])
        self.classifier = Classifier(configuration["classifier"], root=root_folder)

    def close(self):
        pass

    def execute(self, storage: MediaStorage, file: MediaFile):

        im = self.read_image(file.path)
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        bw_im = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)

        detections = self.detector(bw_im)

        labels = {ml.reference: ml for ml in storage.labels()}

        for bbox, info in detections:
            if not info.selected:
                continue
            x0, y0, x1, y1 = bbox
            cls_ref = self.classifier(bbox.crop(im, enlarge=True))
            label = labels.get(cls_ref, cls_ref)
            file.add_bounding_box(x0, y0, bbox.w, bbox.h, label=label)

    def read_image(self, path: str, mode: int = cv2.IMREAD_COLOR) -> np.ndarray:
        return cv2.imread(path, mode)