PyCS/models/moth_scanner/scanner/__init__.py

from typing import List

import cv2
import numpy as np

from json import dump, load

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

from .detector import Detector
from .classifier import Classifier
from .detector import BBox

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 pure_inference(self, storage: MediaStorage, file: MediaFile, bounding_boxes: List[MediaBoundingBox]):
        im = self.read_image(file.path)
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        bw_im = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)

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

        bbox_labels = []
        for bbox in bounding_boxes:
            bbox = BBox(bbox.x, bbox.y, bbox.x + bbox.w , bbox.y + bbox.h)
            x0, y0, x1, y1 = bbox
            cls_ref = self.classifier(bbox.crop(im, enlarge=True))
            label = labels.get(cls_ref, cls_ref)
            bbox_labels.append(label)

        return bbox_labels

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