Update app.py

app.py

@@ -1,44 +1,74 @@
-import io
 import gradio as gr
-from PIL import Image
-import vtracer
+from PIL import Image, ImageFilter
 import numpy as np
-import cv2
+import io
 import tempfile
+import vtracer
+from skimage import feature, filters, morphology
+import cv2
 from rembg import remove
-from skimage import filters, img_as_ubyte, color, measure
-from skimage.restoration import denoise_bilateral
-from sklearn.cluster import KMeans
 
-def convert_image(image, color_mode, hierarchical, mode, filter_speckle,
-                  color_precision, layer_difference, corner_threshold,
-                  length_threshold, max_iterations, splice_threshold, path_precision,
-                  detail_level, ai_edge_enhance, smoothness, background_remove):
-    """Converts an image to SVG with customizable parameters and additional options."""
+def preprocess_image(image, blur_radius, sharpen_radius, noise_reduction, detail_level, edge_method, color_quantization, enhance_with_ai, remove_bg):
+    """Advanced preprocessing of the image before vectorization."""
+    
+    if blur_radius > 0:
+        image = image.filter(ImageFilter.GaussianBlur(blur_radius))
+    
+    if sharpen_radius > 0:
+        image = image.filter(ImageFilter.UnsharpMask(radius=sharpen_radius, percent=150, threshold=3))
+    
+    if noise_reduction > 0:
+        image_np = np.array(image)
+        image_np = cv2.fastNlMeansDenoisingColored(image_np, None, h=noise_reduction, templateWindowSize=7, searchWindowSize=21)
+        image = Image.fromarray(image_np)
 
-    # Optionally remove the background
-    if background_remove:
-        image = Image.open(io.BytesIO(remove(np.array(image))))
+    if detail_level > 0:
+        sigma = max(0.5, 3.0 - (detail_level * 0.5))
+        image_np = np.array(image.convert('L'))
+        
+        if edge_method == 'Canny':
+            edges = feature.canny(image_np, sigma=sigma)
+        elif edge_method == 'Sobel':
+            edges = filters.sobel(image_np)
+        elif edge_method == 'Scharr':
+            edges = filters.scharr(image_np)
+        else:  # Prewitt
+            edges = filters.prewitt(image_np)
+        
+        edges = morphology.dilation(edges, morphology.square(max(1, 6 - detail_level)))
+        edges_img = Image.fromarray((edges * 255).astype(np.uint8))
+        image = Image.blend(image.convert('RGB'), edges_img.convert('RGB'), alpha=0.5)
+    
+    if color_quantization > 0:
+        image = quantize_colors(image, color_quantization)
 
-    # Apply AI enhancement and smoothness if selected
-    if ai_edge_enhance:
+    if enhance_with_ai:
         image_np = np.array(image)
-        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
-        edges = filters.sobel(image_np)
-        enhanced = img_as_ubyte(edges)
-        enhanced = denoise_bilateral(enhanced, sigma_color=0.05, sigma_spatial=15, multichannel=False)
-        image = Image.fromarray(cv2.cvtColor(enhanced, cv2.COLOR_GRAY2RGB))
-    
-    if smoothness > 0:
+        # AI-based enhancement for smoothing edges and improving vectorization
+        image_np = remove(image_np)
+        image = Image.fromarray(image_np)
+
+    if remove_bg:
         image_np = np.array(image)
-        image_np = cv2.GaussianBlur(image_np, (smoothness*2+1, smoothness*2+1), 0)
+        image_np = remove(image_np)
         image = Image.fromarray(image_np)
+    
+    return image
 
+def convert_image(image, blur_radius, sharpen_radius, noise_reduction, detail_level, edge_method, color_quantization,
+                  color_mode, hierarchical, mode, filter_speckle, color_precision, layer_difference, 
+                  corner_threshold, length_threshold, max_iterations, splice_threshold, path_precision,
+                  enhance_with_ai, remove_bg):
+    """Convert an image to SVG using vtracer with customizable and advanced parameters."""
+    
+    # Preprocess the image with additional detail level settings
+    image = preprocess_image(image, blur_radius, sharpen_radius, noise_reduction, detail_level, edge_method, color_quantization, enhance_with_ai, remove_bg)
+    
     # Convert Gradio image to bytes for vtracer compatibility
     img_byte_array = io.BytesIO()
     image.save(img_byte_array, format='PNG')
     img_bytes = img_byte_array.getvalue()
-
+    
     # Perform the conversion
     svg_str = vtracer.convert_raw_image_to_svg(
         img_bytes,
@@ -53,48 +83,69 @@ def convert_image(image, color_mode, hierarchical, mode, filter_speckle,
         length_threshold=float(length_threshold),
         max_iterations=int(max_iterations),
         splice_threshold=int(splice_threshold),
-        path_precision=int(path_precision),
-        detail_level=int(detail_level)
+        path_precision=int(path_precision)
     )
-
+    
     # Save the SVG string to a temporary file
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.svg')
     temp_file.write(svg_str.encode('utf-8'))
     temp_file.close()
-
-    return gr.HTML(f'<svg viewBox="0 0 {image.width} {image.height}">{svg_str}</svg>'), temp_file.name
+    
+    # Display the SVG in the Gradio interface and provide the download link
+    svg_html = f'<svg viewBox="0 0 {image.width} {image.height}">{svg_str}</svg>'
+    return gr.HTML(svg_html), temp_file.name
 
 # Gradio interface
 iface = gr.Blocks()
 
 with iface:
+    gr.Markdown("# Super-Advanced Image to SVG Converter")
     
-    gr.Interface(
-    fn=convert_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Radio(choices=["Color", "Binary"], value="Color", label="Color Mode"),
-        gr.Radio(choices=["Stacked", "Cutout"], value="Stacked", label="Hierarchical"),
-        gr.Radio(choices=["Spline", "Polygon", "None"], value="Spline", label="Mode"),
-        gr.Slider(minimum=1, maximum=100, value=4, step=1, label="Filter Speckle"),
-        gr.Slider(minimum=1, maximum=100, value=6, step=1, label="Color Precision"),
-        gr.Slider(minimum=1, maximum=100, value=16, step=1, label="Layer Difference"),
-        gr.Slider(minimum=1, maximum=100, value=60, step=1, label="Corner Threshold"),
-        gr.Slider(minimum=1, maximum=100, value=4.0, step=0.5, label="Length Threshold"),
-        gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Max Iterations"),
-        gr.Slider(minimum=1, maximum=100, value=45, step=1, label="Splice Threshold"),
-        gr.Slider(minimum=1, maximum=100, value=8, step=1, label="Path Precision"),
-        gr.Slider(minimum=1, maximum=5, value=3, step=1, label="Detail Level"),
-        gr.Checkbox(value=False, label="AI Edge Enhance"),
-        gr.Slider(minimum=0, maximum=10, value=0, step=1, label="Smoothness"),
-        gr.Checkbox(value=False, label="Remove Background")
-    ],
-    outputs=[
-        gr.HTML(label="SVG Output"),
-        gr.File(label="Download SVG")
-    ],
-    title="Advanced Image to SVG Converter",
-    description="Upload an image and customize the conversion parameters as needed.",
-)
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload Image")
+        blur_radius_input = gr.Slider(minimum=0, maximum=10, value=0, step=0.1, label="Blur Radius (for smoothing)")
+        sharpen_radius_input = gr.Slider(minimum=0, maximum=5, value=0, step=0.1, label="Sharpen Radius")
+        noise_reduction_input = gr.Slider(minimum=0, maximum=30, value=0, step=1, label="Noise Reduction")
+        enhance_with_ai_input = gr.Checkbox(label="AI Edge Enhance", value=False)
+        remove_bg_input = gr.Checkbox(label="Remove Background", value=False)
+    
+    with gr.Row():
+        detail_level_input = gr.Slider(minimum=0, maximum=10, value=5, step=1, label="Detail Level")
+        edge_method_input = gr.Radio(choices=["Canny", "Sobel", "Scharr", "Prewitt"], value="Canny", label="Edge Detection Method")
+        color_quantization_input = gr.Slider(minimum=2, maximum=64, value=0, step=2, label="Color Quantization (0 to disable)")
+    
+    with gr.Row():
+        color_mode_input = gr.Radio(choices=["Color", "Binary"], value="Color", label="Color Mode")
+        hierarchical_input = gr.Radio(choices=["Stacked", "Cutout"], value="Stacked", label="Hierarchical")
+        mode_input = gr.Radio(choices=["Spline", "Polygon", "None"], value="Spline", label="Mode")
+    
+    with gr.Row():
+        filter_speckle_input = gr.Slider(minimum=1, maximum=100, value=4, step=1, label="Filter Speckle")
+        color_precision_input = gr.Slider(minimum=1, maximum=100, value=6, step=1, label="Color Precision")
+        layer_difference_input = gr.Slider(minimum=1, maximum=100, value=16, step=1, label="Layer Difference")
+    
+    with gr.Row():
+        corner_threshold_input = gr.Slider(minimum=1, maximum=100, value=60, step=1, label="Corner Threshold")
+        length_threshold_input = gr.Slider(minimum=1, maximum=100, value=4.0, step=0.5, label="Length Threshold")
+        max_iterations_input = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Max Iterations")
+    
+    with gr.Row():
+        splice_threshold_input = gr.Slider(minimum=1, maximum=100, value=45, step=1, label="Splice Threshold")
+        path_precision_input = gr.Slider(minimum=1, maximum=100, value=8, step=1, label="Path Precision")
+    
+    convert_button = gr.Button("Convert Image to SVG")
+    svg_output = gr.HTML(label="SVG Output")
+    download_output = gr.File(label="Download SVG")
+    
+    convert_button.click(
+        fn=convert_image,
+        inputs=[
+            image_input, blur_radius_input, sharpen_radius_input, noise_reduction_input, detail_level_input, edge_method_input, color_quantization_input,
+            color_mode_input, hierarchical_input, mode_input, filter_speckle_input, color_precision_input,
+            layer_difference_input, corner_threshold_input, length_threshold_input, max_iterations_input,
+            splice_threshold_input, path_precision_input, enhance_with_ai_input, remove_bg_input
+        ],
+        outputs=[svg_output, download_output]
+    )
 
 iface.launch()