847 lines
34 KiB
Python
847 lines
34 KiB
Python
"""
|
|
Key to STL Generator - Educational Tool
|
|
Analyzes key images and creates 3D STL models of the key teeth profile
|
|
Enhanced version with manual selection tools
|
|
"""
|
|
|
|
from PIL import Image, ImageDraw, ImageFilter, ImageTk, ImageEnhance
|
|
import numpy as np
|
|
from scipy import ndimage, signal
|
|
from scipy.ndimage import median_filter, gaussian_filter
|
|
import tkinter as tk
|
|
from tkinter import filedialog, messagebox, ttk
|
|
import struct
|
|
|
|
class KeyToSTL:
|
|
def __init__(self):
|
|
self.original_image = None
|
|
self.image = None
|
|
self.processed_image = None
|
|
self.key_profile = []
|
|
self.key_bounds = None
|
|
self.manual_bounds = None
|
|
self.groove_regions = []
|
|
|
|
def load_image(self, filepath):
|
|
"""Load and preprocess the key image"""
|
|
self.original_image = Image.open(filepath)
|
|
self.original_image = self.original_image.convert('RGB')
|
|
self.image = self.original_image.convert('L')
|
|
print(f"Image loaded: {self.image.size}")
|
|
return self.image
|
|
|
|
def set_manual_bounds(self, x1, y1, x2, y2):
|
|
"""Set manual key bounds from user selection"""
|
|
self.manual_bounds = {
|
|
'left': min(x1, x2),
|
|
'right': max(x1, x2),
|
|
'top': min(y1, y2),
|
|
'bottom': max(y1, y2)
|
|
}
|
|
print(f"Manual bounds set: {self.manual_bounds}")
|
|
|
|
def add_groove_region(self, x1, y1, x2, y2):
|
|
"""Add a groove region that should be emphasized"""
|
|
region = {
|
|
'left': min(x1, x2),
|
|
'right': max(x1, x2),
|
|
'top': min(y1, y2),
|
|
'bottom': max(y1, y2)
|
|
}
|
|
self.groove_regions.append(region)
|
|
print(f"Groove region added: {region}")
|
|
|
|
def clear_groove_regions(self):
|
|
"""Clear all groove regions"""
|
|
self.groove_regions = []
|
|
print("Groove regions cleared")
|
|
|
|
def denoise_image(self, median_size=3, gaussian_sigma=1.5):
|
|
"""Apply strong denoising to handle static/noise in images"""
|
|
if self.image is None:
|
|
raise ValueError("No image loaded")
|
|
|
|
img_array = np.array(self.image)
|
|
denoised = median_filter(img_array, size=median_size)
|
|
denoised = gaussian_filter(denoised, sigma=gaussian_sigma)
|
|
|
|
return denoised
|
|
|
|
def auto_detect_key(self, noise_handling=True):
|
|
"""Automatically detect the key region with noise handling"""
|
|
if self.image is None:
|
|
raise ValueError("No image loaded")
|
|
|
|
if noise_handling:
|
|
img_array = self.denoise_image(median_size=5, gaussian_sigma=2.0)
|
|
else:
|
|
img_array = np.array(self.image)
|
|
|
|
img_pil = Image.fromarray(img_array.astype(np.uint8))
|
|
enhancer = ImageEnhance.Contrast(img_pil)
|
|
enhanced = enhancer.enhance(2.5)
|
|
|
|
img_array = np.array(enhanced)
|
|
img_array = gaussian_filter(img_array, sigma=1.0)
|
|
|
|
sobel_x = ndimage.sobel(img_array, axis=1)
|
|
sobel_y = ndimage.sobel(img_array, axis=0)
|
|
edge_magnitude = np.hypot(sobel_x, sobel_y)
|
|
|
|
edge_magnitude = (edge_magnitude / edge_magnitude.max() * 255).astype(np.uint8)
|
|
|
|
threshold = np.percentile(edge_magnitude[edge_magnitude > 0], 90)
|
|
binary_edges = edge_magnitude > threshold
|
|
|
|
from scipy.ndimage import binary_closing, binary_opening, binary_dilation, binary_erosion
|
|
|
|
binary_edges = binary_opening(binary_edges, structure=np.ones((2, 2)))
|
|
binary_edges = binary_closing(binary_edges, structure=np.ones((3, 3)))
|
|
binary_edges = binary_erosion(binary_edges, structure=np.ones((2, 2)))
|
|
binary_edges = binary_dilation(binary_edges, structure=np.ones((2, 2)))
|
|
|
|
self.processed_image = Image.fromarray((binary_edges * 255).astype(np.uint8))
|
|
return self.processed_image
|
|
|
|
def find_key_bounds(self, margin_percent=0.1):
|
|
"""Find the bounding box of the key with noise-resistant algorithm"""
|
|
if self.processed_image is None:
|
|
raise ValueError("Process image first")
|
|
|
|
# Use manual bounds if available
|
|
if self.manual_bounds:
|
|
self.key_bounds = self.manual_bounds.copy()
|
|
print(f"Using manual bounds: {self.key_bounds}")
|
|
return self.key_bounds
|
|
|
|
edges_array = np.array(self.processed_image)
|
|
height, width = edges_array.shape
|
|
|
|
row_content = np.sum(edges_array, axis=1)
|
|
col_content = np.sum(edges_array, axis=0)
|
|
|
|
row_smooth = gaussian_filter(row_content.astype(float), sigma=height * 0.02)
|
|
col_smooth = gaussian_filter(col_content.astype(float), sigma=width * 0.01)
|
|
|
|
row_threshold = np.percentile(row_smooth[row_smooth > 0], 25)
|
|
col_threshold = np.percentile(col_smooth[col_smooth > 0], 15)
|
|
|
|
rows_with_content = np.where(row_smooth > row_threshold)[0]
|
|
cols_with_content = np.where(col_smooth > col_threshold)[0]
|
|
|
|
if len(rows_with_content) > 0 and len(cols_with_content) > 0:
|
|
row_margin = int(height * margin_percent)
|
|
col_margin = int(width * margin_percent * 0.5)
|
|
|
|
self.key_bounds = {
|
|
'top': max(0, rows_with_content[0] - row_margin),
|
|
'bottom': min(height - 1, rows_with_content[-1] + row_margin),
|
|
'left': max(0, cols_with_content[0] - col_margin),
|
|
'right': min(width - 1, cols_with_content[-1] + col_margin)
|
|
}
|
|
print(f"Key bounds detected: {self.key_bounds}")
|
|
else:
|
|
self.key_bounds = {
|
|
'top': height // 4,
|
|
'bottom': 3 * height // 4,
|
|
'left': width // 10,
|
|
'right': 9 * width // 10
|
|
}
|
|
print("Using fallback bounds")
|
|
|
|
return self.key_bounds
|
|
|
|
def extract_key_profile(self, use_top=True, consensus_window=5):
|
|
"""Extract key profile with noise-resistant consensus approach"""
|
|
if self.processed_image is None or self.key_bounds is None:
|
|
raise ValueError("Process image and find bounds first")
|
|
|
|
edges_array = np.array(self.processed_image)
|
|
bounds = self.key_bounds
|
|
|
|
roi = edges_array[bounds['top']:bounds['bottom'],
|
|
bounds['left']:bounds['right']]
|
|
|
|
height, width = roi.shape
|
|
profile = []
|
|
|
|
for x in range(width):
|
|
start_col = max(0, x - consensus_window // 2)
|
|
end_col = min(width, x + consensus_window // 2 + 1)
|
|
|
|
window = roi[:, start_col:end_col]
|
|
|
|
edge_positions = []
|
|
for col in range(window.shape[1]):
|
|
column = window[:, col]
|
|
edge_pixels = np.where(column > 128)[0]
|
|
|
|
if len(edge_pixels) > 0:
|
|
if use_top:
|
|
edge_positions.append(edge_pixels[0])
|
|
else:
|
|
edge_positions.append(edge_pixels[-1])
|
|
|
|
if len(edge_positions) > 0:
|
|
y_pos = int(np.median(edge_positions))
|
|
actual_y = y_pos + bounds['top']
|
|
actual_x = x + bounds['left']
|
|
profile.append((actual_x, actual_y))
|
|
else:
|
|
if len(profile) > 0:
|
|
profile.append((x + bounds['left'], profile[-1][1]))
|
|
else:
|
|
profile.append((x + bounds['left'], bounds['top'] + height // 2))
|
|
|
|
profile = self._fill_profile_gaps(profile)
|
|
|
|
# Apply groove emphasis
|
|
if self.groove_regions:
|
|
profile = self._emphasize_grooves(profile)
|
|
|
|
self.key_profile = profile
|
|
return profile
|
|
|
|
def _emphasize_grooves(self, profile, emphasis_factor=1.5):
|
|
"""Emphasize groove regions in the profile"""
|
|
emphasized = []
|
|
|
|
for x, y in profile:
|
|
# Check if point is in any groove region
|
|
in_groove = False
|
|
for groove in self.groove_regions:
|
|
if groove['left'] <= x <= groove['right']:
|
|
in_groove = True
|
|
break
|
|
|
|
if in_groove:
|
|
# Find average Y in this region
|
|
region_ys = [py for px, py in profile if groove['left'] <= px <= groove['right']]
|
|
if region_ys:
|
|
avg_y = np.mean(region_ys)
|
|
# Emphasize deviation from average
|
|
deviation = y - avg_y
|
|
new_y = avg_y + deviation * emphasis_factor
|
|
emphasized.append((x, int(new_y)))
|
|
else:
|
|
emphasized.append((x, y))
|
|
else:
|
|
emphasized.append((x, y))
|
|
|
|
return emphasized
|
|
|
|
def _fill_profile_gaps(self, profile, max_gap=10):
|
|
"""Fill gaps in profile with interpolation"""
|
|
if len(profile) < 2:
|
|
return profile
|
|
|
|
filled = [profile[0]]
|
|
|
|
for i in range(1, len(profile)):
|
|
prev_x, prev_y = filled[-1]
|
|
curr_x, curr_y = profile[i]
|
|
|
|
if abs(curr_y - prev_y) > max_gap:
|
|
steps = int(abs(curr_x - prev_x))
|
|
if steps > 1:
|
|
for j in range(1, steps):
|
|
interp_x = prev_x + j
|
|
interp_y = int(prev_y + (curr_y - prev_y) * j / steps)
|
|
filled.append((interp_x, interp_y))
|
|
|
|
filled.append((curr_x, curr_y))
|
|
|
|
return filled
|
|
|
|
def smooth_profile_advanced(self, window_size=11, poly_order=3):
|
|
"""Apply Savitzky-Golay filter with noise handling"""
|
|
if not self.key_profile:
|
|
raise ValueError("Extract profile first")
|
|
|
|
profile_array = np.array(self.key_profile)
|
|
x_coords = profile_array[:, 0]
|
|
y_coords = profile_array[:, 1]
|
|
|
|
y_median = median_filter(y_coords, size=5)
|
|
|
|
window_size = min(window_size, len(y_median))
|
|
if window_size % 2 == 0:
|
|
window_size -= 1
|
|
window_size = max(window_size, poly_order + 2)
|
|
if window_size % 2 == 0:
|
|
window_size += 1
|
|
|
|
try:
|
|
smoothed_y = signal.savgol_filter(y_median, window_size, poly_order)
|
|
except:
|
|
smoothed_y = gaussian_filter(y_median, sigma=window_size / 3.0)
|
|
|
|
self.key_profile = [(int(x), int(y)) for x, y in zip(x_coords, smoothed_y)]
|
|
return self.key_profile
|
|
|
|
def remove_outliers(self, threshold=3.0, window_size=10):
|
|
"""Remove outlier points using local statistics"""
|
|
if not self.key_profile:
|
|
return
|
|
|
|
profile_array = np.array(self.key_profile)
|
|
y_coords = profile_array[:, 1].astype(float)
|
|
|
|
cleaned_y = y_coords.copy()
|
|
|
|
for i in range(len(y_coords)):
|
|
start = max(0, i - window_size // 2)
|
|
end = min(len(y_coords), i + window_size // 2 + 1)
|
|
window = y_coords[start:end]
|
|
|
|
local_median = np.median(window)
|
|
mad = np.median(np.abs(window - local_median))
|
|
|
|
if mad > 0:
|
|
z_score = abs(y_coords[i] - local_median) / (1.4826 * mad)
|
|
if z_score > threshold:
|
|
cleaned_y[i] = local_median
|
|
|
|
self.key_profile = [(int(x), int(y)) for x, y in zip(profile_array[:, 0], cleaned_y)]
|
|
return self.key_profile
|
|
|
|
def normalize_profile(self):
|
|
"""Normalize profile to 0-1 range for depth"""
|
|
if not self.key_profile:
|
|
raise ValueError("Extract profile first")
|
|
|
|
profile_array = np.array(self.key_profile)
|
|
y_coords = profile_array[:, 1]
|
|
|
|
y_min, y_max = y_coords.min(), y_coords.max()
|
|
|
|
if y_max - y_min < 1:
|
|
return [(x, 0.5) for x, _ in self.key_profile]
|
|
|
|
normalized = []
|
|
for x, y in self.key_profile:
|
|
normalized_depth = (y - y_min) / (y_max - y_min)
|
|
normalized.append((x, normalized_depth))
|
|
|
|
return normalized
|
|
|
|
def generate_stl(self, output_path, depth_scale=2.0, base_thickness=1.0):
|
|
"""Generate STL file from the key profile"""
|
|
if not self.key_profile:
|
|
raise ValueError("Extract profile first")
|
|
|
|
normalized_profile = self.normalize_profile()
|
|
|
|
profile_array = np.array(self.key_profile)
|
|
x_coords = profile_array[:, 0]
|
|
x_min, x_max = x_coords.min(), x_coords.max()
|
|
|
|
key_length = float(x_max - x_min)
|
|
key_width = 10.0
|
|
|
|
scaled_profile = []
|
|
for (x, _), (_, normalized_depth) in zip(self.key_profile, normalized_profile):
|
|
scaled_x = float(x - x_min)
|
|
actual_depth = normalized_depth * depth_scale
|
|
scaled_profile.append((scaled_x, actual_depth))
|
|
|
|
vertices, faces = self._create_mesh(scaled_profile, key_width, base_thickness)
|
|
self._write_stl(output_path, vertices, faces)
|
|
|
|
print(f"STL file saved: {output_path}")
|
|
print(f"Dimensions: {key_length:.1f} x {key_width:.1f} x {base_thickness + depth_scale:.1f} units")
|
|
|
|
def _create_mesh(self, profile, width, base_thickness):
|
|
"""Create 3D mesh from 2D profile"""
|
|
vertices = []
|
|
faces = []
|
|
|
|
n = len(profile)
|
|
|
|
for x, depth in profile:
|
|
z_top = base_thickness + depth
|
|
vertices.append([x, 0.0, z_top])
|
|
vertices.append([x, width, z_top])
|
|
vertices.append([x, 0.0, 0.0])
|
|
vertices.append([x, width, 0.0])
|
|
|
|
for i in range(n - 1):
|
|
base = i * 4
|
|
next_base = (i + 1) * 4
|
|
|
|
v1, v2 = base, base + 1
|
|
v3, v4 = next_base, next_base + 1
|
|
faces.append([v1, v3, v2])
|
|
faces.append([v2, v3, v4])
|
|
|
|
v1, v2 = base + 2, base + 3
|
|
v3, v4 = next_base + 2, next_base + 3
|
|
faces.append([v1, v2, v3])
|
|
faces.append([v2, v4, v3])
|
|
|
|
v1, v2 = base, base + 2
|
|
v3, v4 = next_base, next_base + 2
|
|
faces.append([v1, v3, v2])
|
|
faces.append([v2, v3, v4])
|
|
|
|
v1, v2 = base + 1, base + 3
|
|
v3, v4 = next_base + 1, next_base + 3
|
|
faces.append([v1, v2, v3])
|
|
faces.append([v2, v4, v3])
|
|
|
|
faces.append([0, 2, 1])
|
|
faces.append([1, 2, 3])
|
|
|
|
last = (n - 1) * 4
|
|
faces.append([last, last + 1, last + 2])
|
|
faces.append([last + 1, last + 3, last + 2])
|
|
|
|
return vertices, faces
|
|
|
|
def _write_stl(self, filepath, vertices, faces):
|
|
"""Write vertices and faces to binary STL file"""
|
|
with open(filepath, 'wb') as f:
|
|
header = b'Key STL - Python Generated' + b' ' * (80 - 26)
|
|
f.write(header)
|
|
f.write(struct.pack('<I', len(faces)))
|
|
|
|
for face in faces:
|
|
v1 = vertices[face[0]]
|
|
v2 = vertices[face[1]]
|
|
v3 = vertices[face[2]]
|
|
|
|
normal = self._calculate_normal(v1, v2, v3)
|
|
|
|
f.write(struct.pack('<fff', *normal))
|
|
f.write(struct.pack('<fff', *v1))
|
|
f.write(struct.pack('<fff', *v2))
|
|
f.write(struct.pack('<fff', *v3))
|
|
f.write(struct.pack('<H', 0))
|
|
|
|
def _calculate_normal(self, v1, v2, v3):
|
|
"""Calculate face normal vector"""
|
|
v1, v2, v3 = np.array(v1), np.array(v2), np.array(v3)
|
|
edge1 = v2 - v1
|
|
edge2 = v3 - v1
|
|
normal = np.cross(edge1, edge2)
|
|
|
|
length = np.linalg.norm(normal)
|
|
if length > 0:
|
|
normal = normal / length
|
|
else:
|
|
normal = np.array([0.0, 0.0, 1.0])
|
|
|
|
return normal.tolist()
|
|
|
|
|
|
class KeySTLGUI:
|
|
def __init__(self, root):
|
|
self.root = root
|
|
self.root.title("Key to STL Generator - Manual Selection Tools")
|
|
self.root.geometry("1200x800")
|
|
|
|
self.key_processor = KeyToSTL()
|
|
self.canvas_image = None
|
|
self.photo_image = None
|
|
|
|
# Selection state
|
|
self.selection_mode = None # 'key_bounds' or 'groove'
|
|
self.selection_start = None
|
|
self.current_rect = None
|
|
self.drawn_rectangles = []
|
|
|
|
self._create_widgets()
|
|
self._bind_canvas_events()
|
|
|
|
def _create_widgets(self):
|
|
main_container = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
|
|
main_container.pack(fill=tk.BOTH, expand=True)
|
|
|
|
# Control panel
|
|
control_frame = ttk.Frame(main_container, padding="10")
|
|
main_container.add(control_frame, weight=0)
|
|
|
|
ttk.Label(control_frame, text="Key to STL Generator",
|
|
font=('Arial', 14, 'bold')).pack(pady=10)
|
|
|
|
# Step 1
|
|
ttk.Label(control_frame, text="Step 1: Load Image",
|
|
font=('Arial', 10, 'bold')).pack(pady=(10, 5))
|
|
ttk.Button(control_frame, text="📁 Load Image",
|
|
command=self.load_image, width=22).pack(pady=5)
|
|
|
|
ttk.Separator(control_frame, orient='horizontal').pack(fill='x', pady=10)
|
|
|
|
# Step 2 - Manual Selection Tools
|
|
ttk.Label(control_frame, text="Step 2: Manual Selection (Optional)",
|
|
font=('Arial', 10, 'bold')).pack(pady=(5, 5))
|
|
|
|
tool_frame = ttk.LabelFrame(control_frame, text="Selection Tools", padding="5")
|
|
tool_frame.pack(fill='x', pady=5)
|
|
|
|
ttk.Button(tool_frame, text="🔲 Select Key Area",
|
|
command=self.start_key_selection, width=20).pack(pady=2)
|
|
ttk.Label(tool_frame, text="Drag rectangle around key",
|
|
font=('Arial', 8), foreground='gray').pack()
|
|
|
|
ttk.Separator(tool_frame, orient='horizontal').pack(fill='x', pady=5)
|
|
|
|
ttk.Button(tool_frame, text="🎯 Mark Groove Region",
|
|
command=self.start_groove_selection, width=20).pack(pady=2)
|
|
ttk.Label(tool_frame, text="Mark areas with teeth cuts",
|
|
font=('Arial', 8), foreground='gray').pack()
|
|
|
|
ttk.Button(tool_frame, text="🗑️ Clear Grooves",
|
|
command=self.clear_grooves, width=20).pack(pady=2)
|
|
|
|
ttk.Button(tool_frame, text="↩️ Cancel Selection",
|
|
command=self.cancel_selection, width=20).pack(pady=2)
|
|
|
|
self.selection_status = ttk.Label(tool_frame, text="No active tool",
|
|
font=('Arial', 8, 'italic'), foreground='blue')
|
|
self.selection_status.pack(pady=5)
|
|
|
|
ttk.Separator(control_frame, orient='horizontal').pack(fill='x', pady=10)
|
|
|
|
# Step 3 - Auto Process
|
|
ttk.Label(control_frame, text="Step 3: Auto Process",
|
|
font=('Arial', 10, 'bold')).pack(pady=(5, 5))
|
|
|
|
self.noise_handling_var = tk.BooleanVar(value=True)
|
|
ttk.Checkbutton(control_frame, text="Enable noise reduction",
|
|
variable=self.noise_handling_var).pack(pady=2)
|
|
|
|
ttk.Button(control_frame, text="🔍 Auto Detect Key",
|
|
command=self.auto_process, width=22).pack(pady=5)
|
|
|
|
self.edge_top_var = tk.BooleanVar(value=True)
|
|
ttk.Checkbutton(control_frame, text="Use top edge",
|
|
variable=self.edge_top_var).pack(pady=2)
|
|
|
|
ttk.Separator(control_frame, orient='horizontal').pack(fill='x', pady=10)
|
|
|
|
# Step 4 - Refine
|
|
ttk.Label(control_frame, text="Step 4: Refine Profile",
|
|
font=('Arial', 10, 'bold')).pack(pady=(5, 5))
|
|
|
|
ttk.Label(control_frame, text="Smoothing:").pack()
|
|
self.smooth_var = tk.IntVar(value=11)
|
|
ttk.Scale(control_frame, from_=5, to=21, variable=self.smooth_var,
|
|
orient=tk.HORIZONTAL).pack()
|
|
|
|
ttk.Button(control_frame, text="✨ Smooth",
|
|
command=self.smooth_profile, width=22).pack(pady=2)
|
|
ttk.Button(control_frame, text="🧹 Remove Outliers",
|
|
command=self.remove_outliers, width=22).pack(pady=2)
|
|
|
|
ttk.Separator(control_frame, orient='horizontal').pack(fill='x', pady=10)
|
|
|
|
# Step 5 - Export
|
|
ttk.Label(control_frame, text="Step 5: Export STL",
|
|
font=('Arial', 10, 'bold')).pack(pady=(5, 5))
|
|
|
|
ttk.Label(control_frame, text="Teeth Depth (mm):").pack()
|
|
self.depth_var = tk.DoubleVar(value=2.0)
|
|
ttk.Entry(control_frame, textvariable=self.depth_var, width=15).pack()
|
|
|
|
ttk.Label(control_frame, text="Base (mm):").pack()
|
|
self.base_var = tk.DoubleVar(value=1.0)
|
|
ttk.Entry(control_frame, textvariable=self.base_var, width=15).pack()
|
|
|
|
ttk.Button(control_frame, text="💾 Generate STL",
|
|
command=self.generate_stl, width=22).pack(pady=10)
|
|
|
|
# Canvas
|
|
canvas_frame = ttk.Frame(main_container)
|
|
main_container.add(canvas_frame, weight=1)
|
|
|
|
self.canvas = tk.Canvas(canvas_frame, bg='#f0f0f0', highlightthickness=0, cursor='crosshair')
|
|
|
|
h_scrollbar = ttk.Scrollbar(canvas_frame, orient=tk.HORIZONTAL, command=self.canvas.xview)
|
|
v_scrollbar = ttk.Scrollbar(canvas_frame, orient=tk.VERTICAL, command=self.canvas.yview)
|
|
|
|
self.canvas.configure(xscrollcommand=h_scrollbar.set, yscrollcommand=v_scrollbar.set)
|
|
|
|
h_scrollbar.pack(side=tk.BOTTOM, fill=tk.X)
|
|
v_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
|
|
self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
|
|
|
|
# Status
|
|
self.status_label = ttk.Label(self.root, text="Ready - Load a key image to start",
|
|
relief=tk.SUNKEN, padding=(5, 2))
|
|
self.status_label.pack(side=tk.BOTTOM, fill=tk.X)
|
|
|
|
def _bind_canvas_events(self):
|
|
"""Bind mouse events for manual selection"""
|
|
self.canvas.bind('<Button-1>', self.on_canvas_click)
|
|
self.canvas.bind('<B1-Motion>', self.on_canvas_drag)
|
|
self.canvas.bind('<ButtonRelease-1>', self.on_canvas_release)
|
|
|
|
def start_key_selection(self):
|
|
"""Start key bounds selection mode"""
|
|
if self.photo_image is None:
|
|
messagebox.showwarning("Warning", "Please load an image first")
|
|
return
|
|
|
|
self.selection_mode = 'key_bounds'
|
|
self.selection_status.config(text="✓ Drag rectangle around key area")
|
|
self.canvas.config(cursor='crosshair')
|
|
self.status_label.config(text="SELECTION MODE: Draw rectangle around the key")
|
|
|
|
def start_groove_selection(self):
|
|
"""Start groove region selection mode"""
|
|
if self.photo_image is None:
|
|
messagebox.showwarning("Warning", "Please load an image first")
|
|
return
|
|
|
|
self.selection_mode = 'groove'
|
|
self.selection_status.config(text="✓ Drag rectangle over groove area")
|
|
self.canvas.config(cursor='crosshair')
|
|
self.status_label.config(text="GROOVE MODE: Draw rectangles over areas with teeth cuts")
|
|
|
|
def cancel_selection(self):
|
|
"""Cancel current selection mode"""
|
|
self.selection_mode = None
|
|
self.selection_start = None
|
|
if self.current_rect:
|
|
self.canvas.delete(self.current_rect)
|
|
self.current_rect = None
|
|
self.selection_status.config(text="No active tool")
|
|
self.canvas.config(cursor='')
|
|
self.status_label.config(text="Selection cancelled")
|
|
|
|
def clear_grooves(self):
|
|
"""Clear all groove regions"""
|
|
self.key_processor.clear_groove_regions()
|
|
# Redraw without groove rectangles
|
|
self._redraw_with_annotations()
|
|
self.status_label.config(text="✓ Groove regions cleared")
|
|
|
|
def on_canvas_click(self, event):
|
|
"""Handle canvas click for selection start"""
|
|
if self.selection_mode:
|
|
self.selection_start = (self.canvas.canvasx(event.x),
|
|
self.canvas.canvasy(event.y))
|
|
if self.current_rect:
|
|
self.canvas.delete(self.current_rect)
|
|
|
|
def on_canvas_drag(self, event):
|
|
"""Handle canvas drag for drawing selection rectangle"""
|
|
if self.selection_mode and self.selection_start:
|
|
x1, y1 = self.selection_start
|
|
x2 = self.canvas.canvasx(event.x)
|
|
y2 = self.canvas.canvasy(event.y)
|
|
|
|
if self.current_rect:
|
|
self.canvas.delete(self.current_rect)
|
|
|
|
# Draw rectangle with different colors for different modes
|
|
color = 'green' if self.selection_mode == 'key_bounds' else 'orange'
|
|
self.current_rect = self.canvas.create_rectangle(
|
|
x1, y1, x2, y2,
|
|
outline=color, width=3, dash=(5, 5)
|
|
)
|
|
|
|
def on_canvas_release(self, event):
|
|
"""Handle canvas release to finalize selection"""
|
|
if self.selection_mode and self.selection_start:
|
|
x1, y1 = self.selection_start
|
|
x2 = self.canvas.canvasx(event.x)
|
|
y2 = self.canvas.canvasy(event.y)
|
|
|
|
# Ensure we have a valid rectangle
|
|
if abs(x2 - x1) > 5 and abs(y2 - y1) > 5:
|
|
if self.selection_mode == 'key_bounds':
|
|
self.key_processor.set_manual_bounds(int(x1), int(y1), int(x2), int(y2))
|
|
self.status_label.config(text="✓ Key area selected")
|
|
# Redraw with permanent rectangle
|
|
if self.current_rect:
|
|
self.canvas.delete(self.current_rect)
|
|
self.current_rect = self.canvas.create_rectangle(
|
|
x1, y1, x2, y2,
|
|
outline='green', width=2
|
|
)
|
|
self.selection_mode = None
|
|
self.selection_status.config(text="Key area set")
|
|
|
|
elif self.selection_mode == 'groove':
|
|
self.key_processor.add_groove_region(int(x1), int(y1), int(x2), int(y2))
|
|
self.status_label.config(text=f"✓ Groove region added (Total: {len(self.key_processor.groove_regions)})")
|
|
# Draw permanent groove rectangle
|
|
if self.current_rect:
|
|
self.canvas.delete(self.current_rect)
|
|
groove_rect = self.canvas.create_rectangle(
|
|
x1, y1, x2, y2,
|
|
outline='orange', width=2, dash=(3, 3)
|
|
)
|
|
self.drawn_rectangles.append(groove_rect)
|
|
# Stay in groove mode for multiple selections
|
|
|
|
self.selection_start = None
|
|
self.current_rect = None
|
|
|
|
def _redraw_with_annotations(self):
|
|
"""Redraw image with all annotations"""
|
|
if self.key_processor.original_image:
|
|
self._display_image(self.key_processor.original_image.convert('L'))
|
|
|
|
# Redraw key bounds if set
|
|
if self.key_processor.manual_bounds:
|
|
bounds = self.key_processor.manual_bounds
|
|
self.canvas.create_rectangle(
|
|
bounds['left'], bounds['top'],
|
|
bounds['right'], bounds['bottom'],
|
|
outline='green', width=2
|
|
)
|
|
|
|
# Redraw groove regions
|
|
self.drawn_rectangles = []
|
|
for groove in self.key_processor.groove_regions:
|
|
rect = self.canvas.create_rectangle(
|
|
groove['left'], groove['top'],
|
|
groove['right'], groove['bottom'],
|
|
outline='orange', width=2, dash=(3, 3)
|
|
)
|
|
self.drawn_rectangles.append(rect)
|
|
|
|
def load_image(self):
|
|
filepath = filedialog.askopenfilename(
|
|
title="Select Key Image",
|
|
filetypes=[("Image files", "*.jpg *.jpeg *.png *.bmp *.gif"), ("All files", "*.*")]
|
|
)
|
|
|
|
if filepath:
|
|
try:
|
|
self.key_processor.load_image(filepath)
|
|
self._display_image(self.key_processor.image)
|
|
self.status_label.config(text=f"✓ Loaded: {filepath.split('/')[-1]}")
|
|
except Exception as e:
|
|
messagebox.showerror("Error", f"Failed to load image: {str(e)}")
|
|
|
|
def auto_process(self):
|
|
try:
|
|
noise_handling = self.noise_handling_var.get()
|
|
|
|
self.status_label.config(text="Processing... applying noise reduction" if noise_handling else "Processing... detecting edges")
|
|
self.root.update()
|
|
|
|
self.key_processor.auto_detect_key(noise_handling=noise_handling)
|
|
|
|
self.status_label.config(text="Processing... finding key bounds")
|
|
self.root.update()
|
|
|
|
self.key_processor.find_key_bounds()
|
|
|
|
self.status_label.config(text="Processing... extracting profile")
|
|
self.root.update()
|
|
|
|
use_top = self.edge_top_var.get()
|
|
profile = self.key_processor.extract_key_profile(use_top=use_top, consensus_window=7)
|
|
|
|
self._display_image_with_profile()
|
|
|
|
groove_info = f" (with {len(self.key_processor.groove_regions)} groove regions)" if self.key_processor.groove_regions else ""
|
|
self.status_label.config(text=f"✓ Profile extracted: {len(profile)} points{groove_info}")
|
|
|
|
except Exception as e:
|
|
messagebox.showerror("Error", f"Auto processing failed: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
def smooth_profile(self):
|
|
try:
|
|
window = self.smooth_var.get()
|
|
if window % 2 == 0:
|
|
window += 1
|
|
self.key_processor.smooth_profile_advanced(window_size=window)
|
|
self._display_image_with_profile()
|
|
self.status_label.config(text="✓ Profile smoothed")
|
|
except Exception as e:
|
|
messagebox.showerror("Error", f"Failed to smooth: {str(e)}")
|
|
|
|
def remove_outliers(self):
|
|
try:
|
|
self.key_processor.remove_outliers()
|
|
self._display_image_with_profile()
|
|
self.status_label.config(text="✓ Outliers removed")
|
|
except Exception as e:
|
|
messagebox.showerror("Error", f"Failed to remove outliers: {str(e)}")
|
|
|
|
def generate_stl(self):
|
|
if not self.key_processor.key_profile:
|
|
messagebox.showwarning("Warning", "Please extract key profile first (Step 3)")
|
|
return
|
|
|
|
filepath = filedialog.asksaveasfilename(
|
|
title="Save STL File",
|
|
defaultextension=".stl",
|
|
filetypes=[("STL files", "*.stl"), ("All files", "*.*")]
|
|
)
|
|
|
|
if filepath:
|
|
try:
|
|
self.key_processor.generate_stl(
|
|
filepath,
|
|
depth_scale=self.depth_var.get(),
|
|
base_thickness=self.base_var.get()
|
|
)
|
|
messagebox.showinfo("Success",
|
|
f"✓ STL file generated!\n\n{filepath}\n\nYou can now scale it in your 3D slicer.")
|
|
self.status_label.config(text=f"✓ STL saved: {filepath.split('/')[-1]}")
|
|
except Exception as e:
|
|
messagebox.showerror("Error", f"Failed to generate STL: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
def _display_image(self, image):
|
|
if image is None:
|
|
return
|
|
|
|
self.photo_image = ImageTk.PhotoImage(image)
|
|
self.canvas.delete("all")
|
|
self.canvas.create_image(0, 0, image=self.photo_image, anchor=tk.NW)
|
|
self.canvas.configure(scrollregion=self.canvas.bbox("all"))
|
|
|
|
def _display_image_with_profile(self):
|
|
if self.key_processor.processed_image and self.key_processor.key_profile:
|
|
display_img = self.key_processor.processed_image.convert('RGB')
|
|
draw = ImageDraw.Draw(display_img)
|
|
|
|
# Draw key bounds
|
|
if self.key_processor.key_bounds:
|
|
bounds = self.key_processor.key_bounds
|
|
draw.rectangle(
|
|
[bounds['left'], bounds['top'], bounds['right'], bounds['bottom']],
|
|
outline='green', width=3
|
|
)
|
|
|
|
# Draw groove regions
|
|
for groove in self.key_processor.groove_regions:
|
|
draw.rectangle(
|
|
[groove['left'], groove['top'], groove['right'], groove['bottom']],
|
|
outline='orange', width=2
|
|
)
|
|
|
|
# Draw profile
|
|
if len(self.key_processor.key_profile) > 1:
|
|
draw.line(self.key_processor.key_profile, fill='red', width=3)
|
|
|
|
for x, y in self.key_processor.key_profile[::10]:
|
|
draw.ellipse([x-2, y-2, x+2, y+2], fill='yellow', outline='red')
|
|
|
|
self._display_image(display_img)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("Key to STL Generator - Manual Selection Tools")
|
|
print("Requirements: pip install scipy pillow numpy")
|
|
print("\nHow to use manual selection:")
|
|
print("1. Load your key image")
|
|
print("2. Click 'Select Key Area' and drag rectangle around key")
|
|
print("3. Click 'Mark Groove Region' and drag rectangles over teeth cuts")
|
|
print("4. Click 'Auto Detect Key' to process with your selections")
|
|
print("5. Refine and export STL")
|
|
|
|
root = tk.Tk()
|
|
app = KeySTLGUI(root)
|
|
root.mainloop() |