200 lines
6.9 KiB
Python
200 lines
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")
|