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

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

    def close(self):
        pass

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

        im = self.read_image(file.path)

        detections = self.detector(im)

        for bbox, info in detections:
            if not info.selected:
                continue
            x0, y0, x1, y1 = bbox
            w, h = x1-x0, y1-y0
            file.add_bounding_box(x0, y0, w, h)

    def read_image(self, path: str) -> np.ndarray:
        im = cv2.imread(path, cv2.IMREAD_COLOR)
        im = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)
        return im