From 676b30db148737dac8d23ccdaed7435fc815573f Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Fri, 4 Jul 2025 08:45:17 -0700 Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=A8=20Inpaint=20tool=20improved=20-=20?= =?UTF-8?q?more=20shapes,=20eraser,=20undos=20(#225)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Anthony Wu --- tools/inpaint_mask_tool.py | 179 +++++++++++++++++++++++++++++++++---- 1 file changed, 160 insertions(+), 19 deletions(-) diff --git a/tools/inpaint_mask_tool.py b/tools/inpaint_mask_tool.py index 07c1a79..c1f8c0d 100644 --- a/tools/inpaint_mask_tool.py +++ b/tools/inpaint_mask_tool.py @@ -15,7 +15,7 @@ class MaskCreator: self.mask_output_path = image_path.with_name(f"{image_path.stem}_mask").with_suffix(".png") # Create a window and set up display image - self.window_name = "MFlux Inpaint Mask Creator - Draw with mouse or trackpad (hot keys: (s)ave, (r)eset, (q)uit" + self.window_name = "MFlux Inpaint Mask Creator - Tools: (b)rush, (e)rase, (r)ectangle, s(q)uare, (o)val, (t)riangle | (u)ndo, (s)ave, (esc) quit" self.display_image = self.original_image.copy() # Create a blank mask the same size as the image @@ -26,6 +26,14 @@ class MaskCreator: self.brush_size = self.BRUSH_SIZES["5"] self.last_point = None + # Tool modes + self.tool_mode = "brush" # brush, erase, rectangle, square, oval, triangle + self.shape_start_point = None + + # Undo history + self.mask_history = [] + self.max_history = 50 + self.overlay = np.zeros_like(self.original_image) # Update display every N drawing events - lower is more responsive @@ -35,12 +43,35 @@ class MaskCreator: # Show the initial display self.update_display() + def save_to_history(self): + self.mask_history.append(self.mask.copy()) + if len(self.mask_history) > self.max_history: + self.mask_history.pop(0) + + def undo(self): + if self.mask_history: + self.mask = self.mask_history.pop() + self.update_display() + print("Undo completed") + else: + print("Nothing to undo") + def mouse_callback(self, event, x, y, flags, param): + if self.tool_mode in ["brush", "erase"]: + self._handle_brush(event, x, y) + else: + self._handle_shape(event, x, y) + + def _handle_brush(self, event, x, y): + # Determine color based on tool mode + color = 0 if self.tool_mode == "erase" else 255 + # Start drawing if event == cv2.EVENT_LBUTTONDOWN: + self.save_to_history() self.drawing = True self.last_point = (x, y) - cv2.circle(self.mask, (x, y), self.brush_size, 255, -1) + cv2.circle(self.mask, (x, y), self.brush_size, color, -1) self.event_counter += 1 if self.event_counter % self.update_frequency == 0: self.update_display() @@ -50,9 +81,9 @@ class MaskCreator: # Use thickness based on brush size for smoother lines if self.last_point: # Ensure we have a last point # Draw a line between the last point and current point - cv2.line(self.mask, self.last_point, (x, y), 255, self.brush_size * 2) + cv2.line(self.mask, self.last_point, (x, y), color, self.brush_size * 2) # Also draw a circle at the current point to avoid gaps in fast movements - cv2.circle(self.mask, (x, y), self.brush_size, 255, -1) + cv2.circle(self.mask, (x, y), self.brush_size, color, -1) self.last_point = (x, y) self.event_counter += 1 if self.event_counter % self.update_frequency == 0: @@ -63,6 +94,83 @@ class MaskCreator: self.drawing = False self.update_display() # Always update display when stopping + def _handle_shape(self, event, x, y): + if event == cv2.EVENT_LBUTTONDOWN: + self.save_to_history() + self.shape_start_point = (x, y) + self.drawing = True + + elif event == cv2.EVENT_MOUSEMOVE and self.drawing: + # Show preview of shape + self.update_display() + self._draw_shape_preview(self.shape_start_point, (x, y)) + cv2.imshow(self.window_name, self.display_image) + + elif event == cv2.EVENT_LBUTTONUP and self.drawing: + self.drawing = False + if self.shape_start_point: + self._draw_shape_on_mask(self.shape_start_point, (x, y)) + self.shape_start_point = None + self.update_display() + + def _get_shape_params(self, start, end): + """Calculate shape parameters based on tool mode.""" + if self.tool_mode == "rectangle": + return {"type": "rectangle", "start": start, "end": end} + elif self.tool_mode == "square": + # Force square aspect ratio + width = abs(end[0] - start[0]) + height = abs(end[1] - start[1]) + size = max(width, height) + # Preserve direction of drag + end_x = start[0] + size if end[0] > start[0] else start[0] - size + end_y = start[1] + size if end[1] > start[1] else start[1] - size + return {"type": "rectangle", "start": start, "end": (end_x, end_y)} + elif self.tool_mode == "oval": + center = ((start[0] + end[0]) // 2, (start[1] + end[1]) // 2) + axes = (abs(end[0] - start[0]) // 2, abs(end[1] - start[1]) // 2) + return {"type": "ellipse", "center": center, "axes": axes} + elif self.tool_mode == "triangle": + return {"type": "polygon", "points": self._get_triangle_points(start, end)} + return None + + def _draw_shape_preview(self, start, end): + overlay = self.display_image.copy() + color = (0, 0, 255) # Red preview + + params = self._get_shape_params(start, end) + if not params: + return + + if params["type"] == "rectangle": + cv2.rectangle(overlay, params["start"], params["end"], color, -1) + elif params["type"] == "ellipse" and params["axes"][0] > 0 and params["axes"][1] > 0: + cv2.ellipse(overlay, params["center"], params["axes"], 0, 0, 360, color, -1) + elif params["type"] == "polygon": + cv2.fillPoly(overlay, [params["points"]], color) + + cv2.addWeighted(overlay, 0.3, self.display_image, 0.7, 0, self.display_image) + + def _draw_shape_on_mask(self, start, end): + params = self._get_shape_params(start, end) + if not params: + return + + if params["type"] == "rectangle": + cv2.rectangle(self.mask, params["start"], params["end"], 255, -1) + elif params["type"] == "ellipse" and params["axes"][0] > 0 and params["axes"][1] > 0: + cv2.ellipse(self.mask, params["center"], params["axes"], 0, 0, 360, 255, -1) + elif params["type"] == "polygon": + cv2.fillPoly(self.mask, [params["points"]], 255) + + def _get_triangle_points(self, start, end): + # Create isosceles triangle + base_center = ((start[0] + end[0]) // 2, end[1]) + apex = (base_center[0], start[1]) + left = (start[0], end[1]) + right = (end[0], end[1]) + return np.array([apex, left, right], np.int32) + def update_display(self): # Create a copy of the original image self.display_image = self.original_image.copy() @@ -75,8 +183,20 @@ class MaskCreator: alpha = 0.5 # Transparency level cv2.addWeighted(self.overlay, alpha, self.display_image, 1 - alpha, 0, self.display_image) - # Draw brush size indicator in the corner - text = f"Brush Size: {self.brush_size} (Hotkeys 1-9: change brush size) " + # Draw tool info in the corner + mode_display = { + "brush": "Brush", + "erase": "Erase", + "rectangle": "Rectangle", + "square": "Square", + "oval": "Oval", + "triangle": "Triangle", + } + text = f"Tool: {mode_display.get(self.tool_mode, self.tool_mode)} | " + if self.tool_mode in ["brush", "erase"]: + text += f"Size: {self.brush_size} (1-9) | " + text += f"History: [{len(self.mask_history)}]" + cv2.putText( self.display_image, text, @@ -96,10 +216,6 @@ class MaskCreator: cv2.imwrite(output_path, self.mask) print(f"Mask saved to {output_path}") - def reset_mask(self): - self.mask = np.zeros(self.original_image.shape[:2], dtype=np.uint8) - self.update_display() - def set_brush_size(self, size_key): if size_key in self.BRUSH_SIZES: self.brush_size = self.BRUSH_SIZES[size_key] @@ -114,21 +230,46 @@ class MaskCreator: key = cv2.waitKey(1) & 0xFF key_char = chr(key) if key < 128 else "" - # Check for brush size hotkeys (1-5) - if key_char in self.BRUSH_SIZES: + # Check for brush size hotkeys (1-9) + if key_char in self.BRUSH_SIZES and self.tool_mode in ["brush", "erase"]: self.set_brush_size(key_char) + # Tool selection + elif key == ord("b"): + self.tool_mode = "brush" + print("Switched to brush tool") + self.update_display() + elif key == ord("e"): + self.tool_mode = "erase" + print("Switched to erase tool") + self.update_display() + elif key == ord("r"): + self.tool_mode = "rectangle" + print("Switched to rectangle tool") + self.update_display() + elif key == ord("q"): + self.tool_mode = "square" + print("Switched to square tool") + self.update_display() + elif key == ord("o"): + self.tool_mode = "oval" + print("Switched to oval tool") + self.update_display() + elif key == ord("t"): + self.tool_mode = "triangle" + print("Switched to triangle tool") + self.update_display() + + # Undo (press 'u') + elif key == ord("u"): + self.undo() + # Save mask (press 's') elif key == ord("s"): self.save_mask(self.mask_output_path) - # Reset mask (press 'r') - elif key == ord("r"): - self.reset_mask() - print("Mask reset") - - # Quit (press 'q' or ESC) - elif key == ord("q") or key == 27: + # Quit (press ESC) + elif key == 27: break cv2.destroyAllWindows()