Update app.py

app.py

@@ -8,8 +8,9 @@ from skimage import color, filters, feature, morphology, exposure, util
 import cv2
 from scipy import ndimage
 from sklearn.cluster import KMeans
+from rembg import remove
 
-def preprocess_image(image, blur_radius, sharpen_radius, noise_reduction, detail_level, edge_method, color_quantization):
+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:
@@ -42,26 +43,23 @@ def preprocess_image(image, blur_radius, sharpen_radius, noise_reduction, detail
     
     if color_quantization > 0:
         image = quantize_colors(image, color_quantization)
+
+    if enhance_with_ai:
+        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)
     
     return image
 
-def quantize_colors(image, n_colors):
-    """Reduces the number of colors in the image using KMeans clustering."""
-    np_image = np.array(image)
-    h, w, c = np_image.shape
-    flat_image = np_image.reshape(-1, c)
-    kmeans = KMeans(n_clusters=n_colors, random_state=42).fit(flat_image)
-    labels = kmeans.predict(flat_image)
-    quantized_image = kmeans.cluster_centers_[labels].reshape(h, w, c).astype('uint8')
-    return Image.fromarray(quantized_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):
+                  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)
+    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()
@@ -99,13 +97,13 @@ iface = gr.Blocks()
 
 with iface:
     gr.Markdown("# Super-Advanced Image to SVG Converter")
-    gr.Markdown("Upload an image and customize the conversion parameters for the highest quality vector results. This tool provides extensive options to analyze and vectorize images at a pixel level with advanced processing techniques.")
     
     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")
@@ -141,10 +139,9 @@ with iface:
             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
+            splice_threshold_input, path_precision_input, enhance_with_ai_input
         ],
         outputs=[svg_output, download_output]
     )
 
-
 iface.launch()