Update app.py

app.py

@@ -1,63 +1,34 @@
 import gradio as gr
-from PIL import Image, ImageFilter, ImageOps
+from PIL import Image, ImageFilter
 import numpy as np
 import io
 import tempfile
 import vtracer
-from skimage import color, filters, exposure
-import torch
-from torchvision import transforms
-
-# AI-based edge detection using a pre-trained PyTorch model (e.g., ResNet)
-class EdgeDetectionModel:
-    def __init__(self, model_name='resnet18', device='cpu'):
-        self.device = device
-        self.model = torch.hub.load('pytorch/vision:v0.10.0', model_name, pretrained=True).to(self.device)
-        self.model.eval()
-
-    def detect_edges(self, image):
-        preprocess = transforms.Compose([
-            transforms.Resize(256),
-            transforms.CenterCrop(224),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-        ])
-        input_tensor = preprocess(image).unsqueeze(0).to(self.device)
-        with torch.no_grad():
-            output = self.model(input_tensor)
-
-        output = output.cpu().numpy()[0].transpose((1, 2, 0))
-        gray_output = color.rgb2gray(output)
-        edges = filters.sobel(gray_output)
-        return edges
-
-# Load the edge detection model
-edge_model = EdgeDetectionModel(device='cpu')
-
-def preprocess_image(image, blur_radius, edge_enhance, contrast_stretch):
+from skimage import color, filters
+from skimage.feature import canny
+from skimage.transform import resize
+
+def preprocess_image(image, blur_radius, edge_enhance, edge_threshold):
     """Applies advanced preprocessing steps to the image before tracing."""
     if blur_radius > 0:
         image = image.filter(ImageFilter.GaussianBlur(blur_radius))
-    
-    if contrast_stretch:
-        # Apply contrast stretching
-        p2, p98 = np.percentile(np.array(image), (2, 98))
-        image = ImageOps.autocontrast(image, cutoff=(p2, p98))
-    
+
     if edge_enhance:
-        edges = edge_model.detect_edges(image)
+        # Convert image to grayscale and apply edge detection
+        gray_image = np.array(image.convert('L'))
+        edges = canny(gray_image, sigma=edge_threshold)
         edges_img = Image.fromarray((edges * 255).astype(np.uint8))
         image = Image.blend(image.convert('RGB'), edges_img.convert('RGB'), alpha=0.5)
-    
+
     return image
 
-def convert_image(image, blur_radius, edge_enhance, contrast_stretch, color_mode, hierarchical, mode, filter_speckle,
-                  color_precision, layer_difference, corner_threshold, length_threshold,
+def convert_image(image, blur_radius, edge_enhance, edge_threshold, color_mode, hierarchical, mode, 
+                  filter_speckle, color_precision, layer_difference, corner_threshold, length_threshold,
                   max_iterations, splice_threshold, path_precision):
-    """Converts an image to SVG using vtracer with customizable parameters and AI enhancements."""
+    """Converts an image to SVG using vtracer with customizable parameters."""
 
     # Preprocess the image
-    image = preprocess_image(image, blur_radius, edge_enhance, contrast_stretch)
+    image = preprocess_image(image, blur_radius, edge_enhance, edge_threshold)
 
     # Convert Gradio image to bytes for vtracer compatibility
     img_byte_array = io.BytesIO()
@@ -94,14 +65,14 @@ def convert_image(image, blur_radius, edge_enhance, contrast_stretch, color_mode
 iface = gr.Blocks()
 
 with iface:
-    gr.Markdown("# Advanced AI-Enhanced Image to SVG Vectors")
-    gr.Markdown("Upload an image and customize the conversion parameters for high-quality vector results. AI-enhanced edge detection, contrast stretching, and other preprocessing options ensure superior vectorization.")
+    gr.Markdown("# CPU-Optimized AI-Enhanced Image to SVG Vectors")
+    gr.Markdown("Upload an image and customize the conversion parameters for high-quality vector results. AI-enhanced edge detection and preprocessing ensure superior vectorization.")
 
     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.5, label="Blur Radius (for smoothing)")
         edge_enhance_input = gr.Checkbox(value=False, label="AI Edge Enhance")
-        contrast_stretch_input = gr.Checkbox(value=False, label="Contrast Stretch")
+        edge_threshold_input = gr.Slider(minimum=0.1, maximum=3.0, value=1.0, step=0.1, label="Edge Detection Threshold")
 
     with gr.Row():
         color_mode_input = gr.Radio(choices=["Color", "Binary"], value="Color", label="Color Mode")
@@ -129,9 +100,10 @@ with iface:
     convert_button.click(
         fn=convert_image,
         inputs=[
-            image_input, blur_radius_input, edge_enhance_input, contrast_stretch_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
+            image_input, blur_radius_input, edge_enhance_input, edge_threshold_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
         ],
         outputs=[svg_output, download_output]
     )