Update app.py

app.py

@@ -7,6 +7,22 @@ import vtracer
 from skimage import feature, filters, morphology
 import cv2
 from rembg import remove
+from sklearn.cluster import KMeans
+
+def quantize_colors(image, num_colors):
+    """Reduce the number of colors in an image."""
+    try:
+        image_np = np.array(image)
+        h, w, c = image_np.shape
+        image_reshaped = image_np.reshape((-1, 3))
+
+        kmeans = KMeans(n_clusters=num_colors, random_state=42).fit(image_reshaped)
+        new_colors = kmeans.cluster_centers_[kmeans.labels_].reshape(h, w, 3).astype(np.uint8)
+        
+        return Image.fromarray(new_colors)
+    except Exception as e:
+        print(f"Error during color quantization: {e}")
+        raise
 
 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."""
@@ -44,13 +60,10 @@ def preprocess_image(image, blur_radius, sharpen_radius, noise_reduction, detail
 
         if enhance_with_ai:
             image_np = np.array(image)
-            # 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 = remove(image_np)
+            # AI-based enhancement for smoothing edges without background removal
+            # Optionally apply background removal only if remove_bg is checked
+            if remove_bg:
+                image_np = remove(image_np)
             image = Image.fromarray(image_np)
 
     except Exception as e:
@@ -62,12 +75,17 @@ def preprocess_image(image, blur_radius, sharpen_radius, noise_reduction, detail
 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):
+                  enhance_with_ai, remove_bg, upscale_factor):
     """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)
 
+    # Upscale the image if needed
+    if upscale_factor > 1:
+        new_size = (int(image.width * upscale_factor), int(image.height * upscale_factor))
+        image = image.resize(new_size, Image.LANCZOS)
+
     # Convert Gradio image to bytes for vtracer compatibility
     img_byte_array = io.BytesIO()
     image.save(img_byte_array, format='PNG')
@@ -117,6 +135,7 @@ with iface:
         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)
+        upscale_factor_input = gr.Slider(minimum=1, maximum=4, value=1, step=0.1, label="Upscale Factor (1 = No Upscaling)")
     
     with gr.Row():
         detail_level_input = gr.Slider(minimum=0, maximum=10, value=5, step=1, label="Detail Level")
@@ -150,11 +169,12 @@ with iface:
         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
+            color_mode_input, hierarchical_input, mode_input, filter_speckle_input, 
+            color_precision_input, length_threshold_input, max_iterations_input, 
+            splice_threshold_input, path_precision_input, enhance_with_ai_input, remove_bg_input, upscale_factor_input
         ],
         outputs=[svg_output, download_output]
     )
 
+
 iface.launch()