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 blob_detector.core.bbox import BBox

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 = None #Classifier(configuration["classifier"], root=root_folder)

    def close(self):
        pass

    def predict(self, im, bboxes, storage):

        for bbox in bboxes:

            if not bbox.is_valid:
                continue

            if self.classifier is None:
                yield bbox, None
                continue

            x0, y0, x1, y1 = bbox
            cls_ref = self.classifier(bbox.crop(im, enlarge=True))
            label = labels.get(cls_ref, cls_ref)

            yield bbox, label


    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)

        bboxes = self.detector(bw_im)

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

        for bbox, label in self.predict(im, bboxes, storage):
            file.add_bounding_box(
                bbox.x0, bbox.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)

        bboxes = [BBox(bbox.x, bbox.y, bbox.x + bbox.w , bbox.y + bbox.h)
            for bbox in bounding_boxes]

        bbox_labels = []
        for bbox, label in self.predict(im, bboxes, storage):
            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)