Update app.py

app.py

@@ -1,71 +1,44 @@
-import gradio as gr
-from PIL import Image, ImageFilter, ImageOps
-import numpy as np
 import io
-import tempfile
+import gradio as gr
+from PIL import Image
 import vtracer
-from skimage import color, filters, feature, morphology, exposure, util
+import numpy as np
 import cv2
-from scipy import ndimage
-from sklearn.cluster import KMeans
+import tempfile
 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 preprocess_image(image, blur_radius, sharpen_radius, noise_reduction, detail_level, edge_method, color_quantization, enhance_with_ai):
-    """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)
+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."""
 
-    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)
+    # Optionally remove the background
+    if background_remove:
+        image = Image.open(io.BytesIO(remove(np.array(image))))
 
-    if enhance_with_ai:
+    # Apply AI enhancement and smoothness if selected
+    if ai_edge_enhance:
         image_np = np.array(image)
-        # AI-based enhancement with rembg for smoothing edges and improving vectorization
-        image_np = remove(image_np)
-        image = Image.fromarray(image_np)
+        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))
     
-    return image
+    if smoothness > 0:
+        image_np = np.array(image)
+        image_np = cv2.GaussianBlur(image_np, (smoothness*2+1, smoothness*2+1), 0)
+        image = Image.fromarray(image_np)
 
-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):
-    """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)
-    
     # 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,
@@ -80,68 +53,48 @@ def convert_image(image, blur_radius, sharpen_radius, noise_reduction, detail_le
         length_threshold=float(length_threshold),
         max_iterations=int(max_iterations),
         splice_threshold=int(splice_threshold),
-        path_precision=int(path_precision)
+        path_precision=int(path_precision),
+        detail_level=int(detail_level)
     )
-    
+
     # 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()
-    
-    # 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
+
+    return gr.HTML(f'<svg viewBox="0 0 {image.width} {image.height}">{svg_str}</svg>'), temp_file.name
 
 # Gradio interface
 iface = gr.Blocks()
 
 with iface:
-    gr.Markdown("# Super-Advanced Image to SVG Converter")
     
-    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)
-    
-    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
-        ],
-        outputs=[svg_output, download_output]
-    )
+    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.",
+)
 
 iface.launch()