ScanSlicer/scanslicer/autoslicer.py

import cv2
import numpy as np

from .autoslicer_detection import AutoslicerDetectionMixin
from .autoslicer_feedback import AutoslicerFeedbackMixin
from .params import AutoslicerParams


# Autoslicer: class definition for ScanSlicer behavior.
class Autoslicer(AutoslicerFeedbackMixin, AutoslicerDetectionMixin):
    MAX_RETURNED_BOXES = 400
    MAX_FEEDBACK_ENTRIES = 500
    AUTOTUNE_PARAM_KEYS = (
        "detect_pipeline",
        "clahe_clip",
        "gaussian",
        "bw_method",
        "bw_thresh",
        "bw_gauss",
        "bw_adaptive_c",
        "invert_mode",
        "morph_op",
        "contour_approx_pct",
        "bbox_min_size",
        "bbox_max_size",
        "bbox_fill_thresh",
        "aspect_ratio_min",
        "aspect_ratio_max",
        "border_ignore",
        "iou_merge_thresh",
        "dilate_kernel",
    )

    # __init__: function definition for ScanSlicer behavior.
    def __init__(self, params=None):
        self.image = None
        self.image_gray = None
        self.image_path = None
        self.abort_flag = False
        self.params = None
        self.set_params(params)
        self.feedback_entries = self.load_feedback_entries()

    # set_params: function definition for ScanSlicer behavior.
    def set_params(self, params):
        self.params = params if params is not None else AutoslicerParams()

    # image_loaded: function definition for ScanSlicer behavior.
    def image_loaded(self):
        return self.image is not None

    # load_image: function definition for ScanSlicer behavior.
    def load_image(self, image_path):
        self.image_path = image_path
        self.image = cv2.imread(image_path)
        self.image_gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)

    # abort_operation: function definition for ScanSlicer behavior.
    def abort_operation(self):
        self.abort_flag = True

    # update_status: function definition for ScanSlicer behavior.
    def update_status(self, update_status_callback, text):
        if update_status_callback is not None:
            update_status_callback(text)

    # save_slice: function definition for ScanSlicer behavior.
    def save_slice(self, hull_quad, out_path, imwrite_params=None):
        hull_quad = np.array(hull_quad, dtype=np.float32).reshape(4, 2)

        # Input order is assumed [top-left, top-right, bottom-right, bottom-left].
        top_w = np.linalg.norm(hull_quad[1] - hull_quad[0])
        bottom_w = np.linalg.norm(hull_quad[2] - hull_quad[3])
        left_h = np.linalg.norm(hull_quad[3] - hull_quad[0])
        right_h = np.linalg.norm(hull_quad[2] - hull_quad[1])

        out_w = max(1, int(round(max(top_w, bottom_w))))
        out_h = max(1, int(round(max(left_h, right_h))))

        dst_quad = np.array(
            [
                [0.0, 0.0],
                [float(out_w - 1), 0.0],
                [float(out_w - 1), float(out_h - 1)],
                [0.0, float(out_h - 1)],
            ],
            dtype=np.float32,
        )

        perspective_m = cv2.getPerspectiveTransform(hull_quad, dst_quad)
        warped_img = cv2.warpPerspective(
            self.image,
            perspective_m,
            (out_w, out_h),
            flags=cv2.INTER_CUBIC,
        )

        if imwrite_params is None:
            cv2.imwrite(out_path, warped_img)
        else:
            cv2.imwrite(out_path, warped_img, imwrite_params)