JunkVision/gui.py

import cv2
import cv2.cv

from threading import Thread

import numpy as np
from functools import partial

class AreaSelector(object):
    def __init__(self, window_name, orig_img, max_areas = 4):
        self.window_name = window_name
        self.orig_img = orig_img
        self.selection = False
        self.area = 0
        self.max_areas = max_areas
        self.rects = []
        for i in xrange(max_areas):
            self.rects.append([])
        self.colors = [
            (0,255,0),
            (255,0,127),
            (0,127,255),
            (255,255,0)
        ]

    @staticmethod
    def handler(event, x, y, flags, self):
        if event == cv2.EVENT_FLAG_RBUTTON:
            self.area += 1
            self.area %= self.max_areas
            if self.selection:
                self.selection = False
                self.redrawRects()
        elif event == cv2.EVENT_LBUTTONDOWN and not self.selection:
            self.selection = True
            self.first_x = x
            self.first_y = y
        elif event == cv2.EVENT_LBUTTONDOWN and self.selection:
            self.selection = False
            self.redrawRects()
        elif event == cv2.EVENT_MOUSEMOVE and self.selection:
            self.rects[self.area] = (min(self.first_x, x), min(self.first_y, y), max(self.first_x, x), max(self.first_y, y))
            self.redrawRects()

    def redrawRects(self):
        cv2.imshow(self.window_name, self.orig_img)
        new_img = self.orig_img.copy()
        for i, r in enumerate(self.rects):
            if not r:
                continue
            cv2.rectangle(new_img, (r[0], r[1]), (r[2], r[3]), self.colors[i], 4)
        cv2.imshow(self.window_name, new_img)

class CorrectionSelector(object):
    def __init__(self, window_name, orig_img):
        self.selecting_rect = True
        self.points = []
        self.cur_point = 0
        self.orig_img = orig_img
        self.window_name = window_name
        self.colors = [
            (0,255,127),
            (255,0,127),
            (0,127,255),
            (255,255,0)
        ]
        self.transform = []
        self.rect_x = 0
        self.rect_y = 0

    @staticmethod
    def handler(event, x, y, flags, self):
        if self.selecting_rect:
            if event == cv2.EVENT_FLAG_RBUTTON:
                self.selecting_rect = False
                self.redraw()
            elif event == cv2.EVENT_FLAG_LBUTTON:
                self.rect_x = x
                self.rect_y = y
            elif event == cv2.EVENT_MOUSEMOVE:
                self.rect = min(self.rect_x, x), min(self.rect_y, y), max(self.rect_x, x), max(self.rect_y, y)
                self.redraw()
        else:
            if event == cv2.EVENT_FLAG_RBUTTON:
                self.redraw()
                self.calc_transform()
                self.selecting_rect = True
            elif event == cv2.EVENT_FLAG_LBUTTON:
                if len(self.points) < 4:
                    self.points.append((x,y))
                else:
                    self.points[self.cur_point] = (x,y)
                self.cur_point += 1
                self.cur_point %= 4
                self.redraw()

    def redraw(self):
        cv2.imshow(self.window_name, self.orig_img)
        new_img = self.orig_img.copy()
        if self.rect:
            cv2.rectangle(new_img, (self.rect[0], self.rect[1]), (self.rect[2], self.rect[3]), (0, 255, 0), 4)
        self.points.sort()
        for i, p in enumerate(self.points):
            cv2.circle(new_img, p, 4, self.colors[i], -1)
        cv2.imshow(self.window_name, new_img)

    def calc_transform(self):
        self.rect_points = [
            (self.rect[0], self.rect[1]),
            (self.rect[0], self.rect[3]),
            (self.rect[2], self.rect[1]),
            (self.rect[2], self.rect[3])
        ]
        sorted_points = []
        best_match = 0
        for i in xrange(4):
            min_delta = cap_width * cap_height
            for j in xrange(4):
                delta = (abs(self.rect_points[i][0] - self.points[j][0]), abs(self.rect_points[i][1] - self.points[j][1]))
                print i, j, delta
                if delta[0] + delta[1] < min_delta:
                    best_match = j
                    min_delta = delta[0] + delta[1]
            print best_match
            sorted_points.append(self.points[best_match])
        print np.array(self.points), np.array(self.rect_points), np.array(sorted_points)

        self.transform = cv2.getPerspectiveTransform(np.array(sorted_points, dtype=np.float32), np.array(self.rect_points, dtype=np.float32))
        print self.transform

    def get_transform(self):
        return self.transform

class ModeSelector(object):
    def __init__(self, gui, min_mode, max_mode):
        self.gui = gui
        self.min_mode = min_mode
        self.max_mode = max_mode
        self.mode = self.gui.mode

    @staticmethod
    def handler(event, x, y, flags, self):
        if event == cv2.EVENT_MBUTTONDOWN:
            mode = self.gui.mode
            mode += 1
            mode %= self.max_mode
            if mode < self.min_mode:
                mode = self.min_mode
            self.gui.mode = mode

class Gui(object):
    def __init__(self, engine):
        self.engine = engine
        self.enabled = False
        self.mode = -2

    def enable(self):
        if self.enabled:
            raise RuntimeError("Gui has already been started!")
        self.thread = Thread(target=self.run, name="JunkVision GUI")
        self.enabled = True
        self.thread.start()

    def set_mode(self, mode):
        self.mode = mode

    def disable(self):
        self.enabled = False
        if self.thread.is_alive:
            self.thread.join()

    def is_enabled(self):
        return self.enabled

    def run(self):
        cv2.namedWindow("JunkVision GUI")
        while self.enabled:
            if self.mode == -2:
                self.selector = AreaSelector("JunkVision GUI")
                cv2.setMouseCallback("JunkVision GUI", CorrectionSelector.handler, self.selector)
                cv2.waitKey()
            elif self.mode == -1:
                self.selector = CorrectionSelector("JunkVision GUI")
                cv2.setMouseCallback("JunkVision GUI", AreaSelector.handler, self.selector)
                cv2.waitKey()
            elif self.mode >= 0 and self.mode <= 2:
                if not isinstance(self.selector, ModeSelector):
                    self.selector = ModeSelector(self, 0, 2)
                    cv2.setMouseCallback("JunkVision GUI", ModeSelector.handler, self.selector)
                if self.mode == 0:
                    img = self.engine.get_plain_image()
                elif self.mode == 1:
                    img = self.engine.get_movement_image()
                elif self.mode == 2:
                    img = self.engine.get_difference_image()

                cv2.imshow("JunkVision GUI", img)
                if cv2.waitKey(10) == 27:
                    self.enabled = False
        cv2.destroyWindow("JunkVision GUI")