HTequi/blob-simulation/detection/limits_maker.py

import cv2


class LimitsMaker:

    def __init__(self, img, scale, window_name, name):
        self.limits = []
        self.max_limits = 4
        self.orig = img
        self.img = img.copy()
        self.scale = scale
        self.window_name = window_name
        self.done = False
        self.limits_drawn = False
        self.name = name

    def add_limit(self, x, y):
        x_img = int(x / self.scale)
        y_img = int(y / self.scale)
        self.limits.append((x_img, y_img))
        cv2.drawMarker(self.img, (x_img, y_img), (0, 0, 255), thickness=5)

    def draw(self):
        if len(self.limits) == self.max_limits and not self.limits_drawn:
            for i, limit in enumerate(self.limits):
                cv2.line(self.img, self.limits[i-1], limit, (0, 0, 255), thickness=3)
            self.limits_drawn = True

        cv2.imshow(self.window_name, cv2.resize(self.img, (0, 0), fx=self.scale, fy=self.scale))

        if self.enough_data():
            self.confirm()

    def enough_data(self):
        return len(self.limits) == self.max_limits

    def compute(self):
        return self.limits

    def toJSON(self):
        return {'Limits': self.limits}

    def help(self):
        print("--- " + self.name + ": Click on the {} corners.".format(self.max_limits))
        print("--- " + self.name + ": Please start from left corner and do it in the right order")

    def on_mouse(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONUP and not self.enough_data():
            if len(self.limits) < self.max_limits:
                self.add_limit(x, y)

    def clear(self):
        self.limits = []
        self.limits_drawn = False
        self.img = self.orig.copy()
        self.done = False

    def confirm(self):
        print("--- " + self.name + ": Press enter if you're ok with data or any other key if you want to restart "
                                   "setup...")
        key = cv2.waitKey(0) & 0xFF
        if key == 13:  # Enter
            print("--- " + self.name + ": " + str(self.compute()))
            self.done = True
        else:
            self.clear()
            self.help()