python Sketcher utility class

This commit is contained in:
Alexander Mordvintsev 2011-07-01 17:33:29 +00:00
parent ee10a38851
commit 84568cb18e
2 changed files with 37 additions and 26 deletions

View File

@ -1,5 +1,5 @@
import numpy as np
import cv2
import cv2, cv
import os
image_extensions = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.pbm', '.pgm', '.ppm']
@ -57,3 +57,29 @@ def mtx2rvec(R):
def draw_str(dst, (x, y), s):
cv2.putText(dst, s, (x+1, y+1), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 0), thickness = 2, linetype=cv2.CV_AA)
cv2.putText(dst, s, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), linetype=cv2.CV_AA)
class Sketcher:
def __init__(self, windowname, dests, colors_func):
self.prev_pt = None
self.windowname = windowname
self.dests = dests
self.colors_func = colors_func
self.dirty = False
self.show()
cv2.setMouseCallback(self.windowname, self.on_mouse)
def show(self):
cv2.imshow(self.windowname, self.dests[0])
def on_mouse(self, event, x, y, flags, param):
pt = (x, y)
if event == cv.CV_EVENT_LBUTTONDOWN:
self.prev_pt = pt
if self.prev_pt and flags & cv.CV_EVENT_FLAG_LBUTTON:
for dst, color in zip(self.dests, self.colors_func()):
cv.Line(dst, self.prev_pt, pt, color, 5)
self.dirty = True
self.prev_pt = pt
self.show()
else:
self.prev_pt = None

View File

@ -1,5 +1,6 @@
import numpy as np
import cv2, cv
from common import Sketcher
help_message = '''
USAGE: watershed.py [<image>]
@ -21,34 +22,17 @@ class App:
self.cur_marker = 1
self.colors = np.int32( list(np.ndindex(2, 2, 2)) ) * 255
cv2.imshow('img', self.markers_vis)
self.prev_pt = None
self.need_update = False
self.auto_update = True
cv2.setMouseCallback('img', self.onmouse)
def onmouse(self, event, x, y, flags, param):
pt = (x, y)
if event == cv.CV_EVENT_LBUTTONDOWN:
self.prev_pt = pt
if self.prev_pt and flags & cv.CV_EVENT_FLAG_LBUTTON:
color = map(int, self.colors[self.cur_marker])
cv.Line(self.markers, self.prev_pt, pt, self.cur_marker, 5)
cv.Line(self.markers_vis, self.prev_pt, pt, color, 5)
self.need_update = True
self.prev_pt = pt
cv2.imshow('img', self.markers_vis)
else:
self.prev_pt = None
self.sketch = Sketcher('img', [self.markers_vis, self.markers], self.get_colors)
def get_colors(self):
return map(int, self.colors[self.cur_marker]), self.cur_marker
def watershed(self):
m = self.markers.copy()
cv2.watershed(self.img, m)
vis = np.uint8( (self.img + self.colors[np.maximum(m, 0)]) / 2 )
cv2.imshow('watershed', vis)
self.need_update = False
def run(self):
while True:
@ -58,15 +42,16 @@ class App:
if ch >= ord('1') and ch <= ord('7'):
self.cur_marker = ch - ord('0')
print 'marker: ', self.cur_marker
if ch == ord(' ') or (self.need_update and self.auto_update):
if ch == ord(' ') or (self.sketch.dirty and self.auto_update):
self.watershed()
self.sketch.dirty = False
if ch in [ord('a'), ord('A')]:
self.auto_update = not self.auto_update
print 'auto_update if', ['off', 'on'][self.auto_update]
if ch in [ord('r'), ord('R')]:
self.markers = np.zeros(self.img.shape[:2], np.int32)
self.markers_vis = self.img.copy()
cv2.imshow('img', self.markers_vis)
self.markers[:] = 0
self.markers_vis[:] = self.img
self.sketch.show()
cv2.destroyWindow('watershed')