
Still WIP, message handler does something at least if started. GUI is somewhat broken though.
200 lines
No EOL
6.9 KiB
Python
200 lines
No EOL
6.9 KiB
Python
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") |