class FFT2DDemo(ImageProcessDemo): YAXIS_DIRECTION = "up" TITLE = "FFT2D Demo" DEFAULT_IMAGE = "lena.jpg" mask_artist = Instance(ImageMaskDrawer) reset = Button(u"Reset") def __init__(self, **kw): super(FFT2DDemo, self).__init__(**kw) self.mask_artist = None def control_panel(self): return VGroup( "reset", show_labels=False, ) def init_draw(self): self.mask_artist = ImageMaskDrawer(self.axe, mask_shape=(512, 512), canmove=False, size=15) self.mask_artist.on_trait_change(self.draw, "mask_updated") def _reset_fired(self): self.mask_artist.clear_mask() self.mask_artist.update() self.draw() def _img_changed(self): w, h = self.img.shape[:2] if w >= 512 and h >= 512: self.img_gray = cv2.cvtColor(self.img, cv2.COLOR_BGR2GRAY)[:512, :512].copy() self.img_freq = fft.fft2(self.img_gray) self.img_mag = fft.fftshift(np.log10(np.abs(self.img_freq))) self.img_show = np.hstack((self.img[:512, :512, :], self.img[:512, :512, :])) self.img_show[:512, 512:, :] = self.img_gray[:512, :512, None] img_uint8 = normalize_gray_image(self.img_mag) self.img_show[:512, :512, :] = img_uint8[:, :, None] def draw(self): if self.mask_artist is None: self.draw_image(self.img_show) return mask = self.mask_artist.get_mask_array().astype(np.bool) if np.any(mask): y, x = np.where(mask) mask[512-y+1, 512-x+1] = True self.mask_artist.array[:, :, -1] = mask * 255 filtered_img = fft.ifft2(self.img_freq * fft.fftshift(mask)).real filtered_img = normalize_gray_image(filtered_img) self.img_show[:512, 512:, :] = filtered_img[:, :, None] else: self.img_show[:512, 512:, :] = self.img_gray[:, :, None] self.draw_image(self.img_show)
class WatershedDemo(ImageProcessDemo): TITLE = "Watershed Demo" DEFAULT_IMAGE = "fruits.jpg" YAXIS_DIRECTION = "up" current_label = Enum(list(range(1, 10))) def __init__(self, **kw): super(WatershedDemo, self).__init__(**kw) self.figure.canvas_events = [("button_press_event", self.on_button_pressed)] self.result_artist = self.axe.imshow(np.zeros((10, 10, 3), np.uint8), alpha=0.3) self.mask_artist = None self.markers = None def control_panel(self): return VGroup(Item("current_label", label="当前标签")) def on_button_pressed(self, event): if event.inaxes is self.axe and event.button == 3: if self.current_label < 9: self.current_label += 1 else: self.current_label = 1 def init_draw(self): h, w = self.img.shape[:2] if self.mask_artist is not None: self.mask_artist.remove() self.mask_artist = ImageMaskDrawer(self.axe, mask_shape=(h, w), canmove=False, size=10) self.mask_artist.on_trait_change(self.update_markers, "drawed") self.markers = np.zeros((h, w), dtype=np.int32) self.result = np.zeros((h, w), dtype=np.int32) def update_markers(self): mask = self.mask_artist.get_mask_array().astype(bool) y, x = np.where(mask) self.markers[y, x] = self.current_label self.mask_artist.clear_mask() self.draw() def draw(self): if self.markers is not None: if self.img.shape[:2] != self.markers.shape: self.init_draw() self.result[:] = self.markers cv2.watershed(self.img, self.result) img = MARKER_COLORS[self.result] self.draw_image(img, self.result_artist, draw=False) self.draw_image(self.img)
class WatershedDemo(ImageProcessDemo): TITLE = u"Watershed Demo" DEFAULT_IMAGE = u"fruits.jpg" YAXIS_DIRECTION = "up" current_label = Enum(range(1, 10)) def __init__(self, **kw): super(WatershedDemo, self).__init__(**kw) self.figure.canvas_events = [ ("button_press_event", self.on_button_pressed) ] self.result_artist = self.axe.imshow(np.zeros((10, 10, 3), np.uint8), alpha=0.3) self.mask_artist = None self.markers = None def control_panel(self): return VGroup( Item("current_label", label=u"当前标签") ) def on_button_pressed(self, event): if event.inaxes is self.axe and event.button == 3: if self.current_label < 9: self.current_label += 1 else: self.current_label = 1 def init_draw(self): h, w = self.img.shape[:2] if self.mask_artist is not None: self.mask_artist.remove() self.mask_artist = ImageMaskDrawer(self.axe, mask_shape=(h, w), canmove=False, size=10) self.mask_artist.on_trait_change(self.update_markers, "drawed") self.markers = np.zeros((h, w), dtype=np.int32) self.result = np.zeros((h, w), dtype=np.int32) def update_markers(self): mask = self.mask_artist.get_mask_array().astype(bool) y, x = np.where(mask) self.markers[y, x] = self.current_label self.mask_artist.clear_mask() self.draw() def draw(self): if self.markers is not None: if self.img.shape[:2] != self.markers.shape: self.init_draw() self.result[:] = self.markers cv2.watershed(self.img, self.result) img = MARKER_COLORS[self.result] self.draw_image(img, self.result_artist, draw=False) self.draw_image(self.img)
class FFT2DDemo(ImageProcessDemo): YAXIS_DIRECTION = "up" TITLE = "FFT2D Demo" DEFAULT_IMAGE = "lena.jpg" mask_artist = Instance(ImageMaskDrawer) reset = Button(u"Reset") def __init__(self, **kw): super(FFT2DDemo, self).__init__(**kw) self.mask_artist = None self.figure.canvas_events = [ ("button_press_event", self.on_button_pressed) ] def control_panel(self): return VGroup( "reset", show_labels=False, ) def init_draw(self): self.mask_artist = ImageMaskDrawer(self.axe, mask_shape=(512, 512), canmove=False, size=15) self.connect_dirty("mask_artist.mask_updated") def on_button_pressed(self, event): if event.inaxes is self.axe: if event.button == 3: self._reset_fired() def _reset_fired(self): self.mask_artist.clear_mask() self.mask_artist.update() self.draw() def _img_changed(self): img = cv2.resize(self.img, (FFT_SIZE, FFT_SIZE)) w, h = img.shape[:2] self.img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:FFT_SIZE, :FFT_SIZE].copy() self.img_freq = fft.fft2(self.img_gray) self.img_mag = fft.fftshift(np.log10(np.abs(self.img_freq))) self.img_show = np.hstack((img[:FFT_SIZE, :FFT_SIZE, :], img[:FFT_SIZE, :FFT_SIZE, :])) self.img_show[:FFT_SIZE, FFT_SIZE:, :] = self.img_gray[:FFT_SIZE, :FFT_SIZE, None] img_uint8 = normalize_gray_image(self.img_mag) self.img_show[:FFT_SIZE, :FFT_SIZE, :] = img_uint8[:, :, None] def draw(self): if self.mask_artist is None: self.draw_image(self.img_show) return mask = self.mask_artist.get_mask_array().astype(np.bool) if np.any(mask): y, x = np.where(mask) m = (y > 0) & (x > 0) x, y = x[m], y[m] mask[FFT_SIZE - y, FFT_SIZE - x] = True self.mask_artist.array[:, :, -1] = mask * 255 filtered_img = fft.ifft2(self.img_freq * fft.fftshift(mask)).real filtered_img = normalize_gray_image(filtered_img) self.img_show[:FFT_SIZE, FFT_SIZE:, :] = filtered_img[:, :, None] else: self.img_show[:FFT_SIZE, FFT_SIZE:, :] = self.img_gray[:, :, None] self.draw_image(self.img_show)
class FFT2DDemo(ImageProcessDemo): YAXIS_DIRECTION = "up" TITLE = "FFT2D Demo" DEFAULT_IMAGE = "lena.jpg" mask_artist = Instance(ImageMaskDrawer) reset = Button(u"Reset") def __init__(self, **kw): super(FFT2DDemo, self).__init__(**kw) self.mask_artist = None self.figure.canvas_events = [("button_press_event", self.on_button_pressed)] def control_panel(self): return VGroup( "reset", show_labels=False, ) def init_draw(self): self.mask_artist = ImageMaskDrawer(self.axe, mask_shape=(512, 512), canmove=False, size=15) self.connect_dirty("mask_artist.mask_updated") def on_button_pressed(self, event): if event.inaxes is self.axe: if event.button == 3: self._reset_fired() def _reset_fired(self): self.mask_artist.clear_mask() self.mask_artist.update() self.draw() def _img_changed(self): img = cv2.resize(self.img, (FFT_SIZE, FFT_SIZE)) w, h = img.shape[:2] self.img_gray = cv2.cvtColor( img, cv2.COLOR_BGR2GRAY)[:FFT_SIZE, :FFT_SIZE].copy() self.img_freq = fft.fft2(self.img_gray) self.img_mag = fft.fftshift(np.log10(np.abs(self.img_freq))) self.img_show = np.hstack( (img[:FFT_SIZE, :FFT_SIZE, :], img[:FFT_SIZE, :FFT_SIZE, :])) self.img_show[:FFT_SIZE, FFT_SIZE:, :] = self.img_gray[:FFT_SIZE, :FFT_SIZE, None] img_uint8 = normalize_gray_image(self.img_mag) self.img_show[:FFT_SIZE, :FFT_SIZE, :] = img_uint8[:, :, None] def draw(self): if self.mask_artist is None: self.draw_image(self.img_show) return mask = self.mask_artist.get_mask_array().astype(np.bool) if np.any(mask): y, x = np.where(mask) m = (y > 0) & (x > 0) x, y = x[m], y[m] mask[FFT_SIZE - y, FFT_SIZE - x] = True self.mask_artist.array[:, :, -1] = mask * 255 filtered_img = fft.ifft2(self.img_freq * fft.fftshift(mask)).real filtered_img = normalize_gray_image(filtered_img) self.img_show[:FFT_SIZE, FFT_SIZE:, :] = filtered_img[:, :, None] else: self.img_show[:FFT_SIZE, FFT_SIZE:, :] = self.img_gray[:, :, None] self.draw_image(self.img_show)
class InPaintDemo(ImageProcessDemo): YAXIS_DIRECTION = "up" TITLE = u"Inpaint Demo" DEFAULT_IMAGE = "stuff.jpg" mask_artist = Instance(ImageMaskDrawer) r = Range(2.0, 20.0, 10.0) # inpaint的半径参数 method = Enum("INPAINT_NS", "INPAINT_TELEA") # inpaint的算法 show_mask = Bool(False) # 是否显示选区 clear_mask = Button(u"清除选区") apply = Button(u"保存结果") def control_panel(self): return VGroup( Item("r", label=u"inpaint半径"), Item("method", label=u"inpaint算法"), Item("show_mask", label=u"显示选区"), Item("clear_mask", show_label=False), Item("apply", show_label=False), ) def __init__(self, **kw): super(InPaintDemo, self).__init__(**kw) self.connect_dirty("r, method") def init_draw(self): self.mask_artist = ImageMaskDrawer(self.axe, self.img, canmove=False, size=15) self.mask_artist.on_trait_change(self.draw, "drawed") def draw(self): if self.mask_artist is None: self.draw_image(self.img) return mask = self.mask_artist.get_mask_array() if self.img.shape[:2] == mask.shape: img2 = cv2.inpaint(self.img, mask, self.r, getattr(cv2, self.method)) self.img2 = img2 self.show_mask = False self.mask_artist.hide_mask() self.draw_image(img2) else: self.draw_image(self.img) def _img_changed(self): if self.mask_artist is not None: self.mask_artist.create_mask(img=self.img) def _show_mask_changed(self): if self.show_mask: self.mask_artist.show_mask() else: self.mask_artist.hide_mask() self.figure.canvas.draw() def _clear_mask_fired(self): self.mask_artist.clear_mask() self.draw() def _apply_fired(self): if hasattr(self, "img2"): self.img[:] = self.img2[:] self.clear_mask = True