296 lines
11 KiB
Python
296 lines
11 KiB
Python
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
class MaskCreator:
|
|
BRUSH_SIZES = {"1": 2, "2": 4, "3": 8, "4": 16, "5": 32, "6": 48, "7": 96, "8": 192, "9": 384}
|
|
|
|
def __init__(self, image_path: Path):
|
|
self.original_image = cv2.imread(image_path)
|
|
if self.original_image is None:
|
|
raise FileNotFoundError(f"Could not open or find the image: {image_path}")
|
|
|
|
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 - 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
|
|
self.mask = np.zeros(self.original_image.shape[:2], dtype=np.uint8)
|
|
|
|
# Set up drawing parameters
|
|
self.drawing = False
|
|
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
|
|
self.update_frequency = 1
|
|
self.event_counter = 0
|
|
|
|
# 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, color, -1)
|
|
self.event_counter += 1
|
|
if self.event_counter % self.update_frequency == 0:
|
|
self.update_display()
|
|
|
|
# Continue drawing
|
|
elif event == cv2.EVENT_MOUSEMOVE and self.drawing:
|
|
# 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), 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, color, -1)
|
|
self.last_point = (x, y)
|
|
self.event_counter += 1
|
|
if self.event_counter % self.update_frequency == 0:
|
|
self.update_display()
|
|
|
|
# Stop drawing
|
|
elif event == cv2.EVENT_LBUTTONUP:
|
|
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()
|
|
|
|
# Create a colored overlay for the mask (semi-transparent red)
|
|
self.overlay[:] = 0 # Reset overlay
|
|
self.overlay[self.mask > 0] = [0, 0, 255] # Red overlay
|
|
|
|
# Apply the overlay
|
|
alpha = 0.5 # Transparency level
|
|
cv2.addWeighted(self.overlay, alpha, self.display_image, 1 - alpha, 0, self.display_image)
|
|
|
|
# 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,
|
|
(10, 30),
|
|
cv2.FONT_HERSHEY_DUPLEX,
|
|
0.5,
|
|
(255, 0, 0),
|
|
1,
|
|
cv2.LINE_AA, # line type: anti-aliased
|
|
)
|
|
|
|
# Display the result with OpenCV's high GUI priority
|
|
cv2.imshow(self.window_name, self.display_image)
|
|
cv2.waitKey(1) # Process events to force display update
|
|
|
|
def save_mask(self, output_path):
|
|
cv2.imwrite(output_path, self.mask)
|
|
print(f"Mask saved to {output_path}")
|
|
|
|
def set_brush_size(self, size_key):
|
|
if size_key in self.BRUSH_SIZES:
|
|
self.brush_size = self.BRUSH_SIZES[size_key]
|
|
print(f"Brush size {size_key}: {self.brush_size}")
|
|
self.update_display()
|
|
|
|
def run(self):
|
|
cv2.namedWindow(self.window_name)
|
|
cv2.setMouseCallback(self.window_name, self.mouse_callback)
|
|
|
|
while True:
|
|
key = cv2.waitKey(1) & 0xFF
|
|
key_char = chr(key) if key < 128 else ""
|
|
|
|
# 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)
|
|
|
|
# Quit (press ESC)
|
|
elif key == 27:
|
|
break
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Create binary mask image from source image to use as the complementary --masked-image-path arg."
|
|
)
|
|
parser.add_argument("image_path", type=Path, help="Path to the input image")
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
mask_creator = MaskCreator(args.image_path)
|
|
mask_creator.run()
|
|
except FileNotFoundError as e:
|
|
print(f"Error: {e}")
|
|
except Exception as e: # noqa
|
|
print(f"An unexpected error occurred: {e}")
|
|
except KeyboardInterrupt:
|
|
pass
|