Update app.py

app.py

@@ -7,99 +7,129 @@ import vtracer
 from skimage import feature, filters, morphology
 import cv2
 from rembg import remove
+import torch
+from transformers import AutoModelForImageSegmentation, AutoProcessor
+import requests
+from huggingface_hub import hf_hub_download
+
+# Load additional Hugging Face models
+segmentation_model = AutoModelForImageSegmentation.from_pretrained("facebook/dino-vitb16")
+segmentation_processor = AutoProcessor.from_pretrained("facebook/dino-vitb16")
 
 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)
-
-    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)
-
-    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)
-        image = Image.fromarray(image_np)
-    
+    try:
+        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)
+
+        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)
+
+        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)
+            image = Image.fromarray(image_np)
+
+    except Exception as e:
+        print(f"Error during preprocessing: {e}")
+        raise
+
     return image
 
+def vectorize_with_hf_model(image):
+    """Vectorizes the image using a Hugging Face model for segmentation or enhancement."""
+    inputs = segmentation_processor(images=image, return_tensors="pt")
+    outputs = segmentation_model(**inputs)
+    mask = outputs["masks"][0][0].cpu().detach().numpy()
+    mask = (mask > 0.5).astype(np.uint8) * 255
+    mask_image = Image.fromarray(mask)
+    return mask_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, 
+                  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, model_choice):
     """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)
-    
+
+    # If a specific model is chosen, use it to process the image before vectorization
+    if model_choice == "Hugging Face Segmentation Model":
+        image = vectorize_with_hf_model(image)
+
     # 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,
-        img_format='png',
-        colormode=color_mode.lower(),
-        hierarchical=hierarchical.lower(),
-        mode=mode.lower(),
-        filter_speckle=int(filter_speckle),
-        color_precision=int(color_precision),
-        layer_difference=int(layer_difference),
-        corner_threshold=int(corner_threshold),
-        length_threshold=float(length_threshold),
-        max_iterations=int(max_iterations),
-        splice_threshold=int(splice_threshold),
-        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()
-    
-    # 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
+
+    try:
+        # Perform the conversion
+        svg_str = vtracer.convert_raw_image_to_svg(
+            img_bytes,
+            img_format='png',
+            colormode=color_mode.lower(),
+            hierarchical=hierarchical.lower(),
+            mode=mode.lower(),
+            filter_speckle=int(filter_speckle),
+            color_precision=int(color_precision),
+            layer_difference=int(layer_difference),
+            corner_threshold=int(corner_threshold),
+            length_threshold=float(length_threshold),
+            max_iterations=int(max_iterations),
+            splice_threshold=int(splice_threshold),
+            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()
+
+        # 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
+
+    except Exception as e:
+        print(f"Error during vectorization: {e}")
+        return f"Error: {e}", None
 
 # Gradio interface
 iface = gr.Blocks()
 
 with iface:
-    gr.Markdown("# Super-Advanced Image to SVG Converter")
+    gr.Markdown("# Super-Advanced Image to SVG Converter with Enhanced Models")
     
     with gr.Row():
         image_input = gr.Image(type="pil", label="Upload Image")
@@ -133,6 +163,8 @@ with iface:
         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")
     
+    model_choice_input = gr.Radio(choices=["None", "Hugging Face Segmentation Model"], value="None", label="Choose Model")
+    
     convert_button = gr.Button("Convert Image to SVG")
     svg_output = gr.HTML(label="SVG Output")
     download_output = gr.File(label="Download SVG")
@@ -143,7 +175,7 @@ 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, enhance_with_ai_input, remove_bg_input
+            splice_threshold_input, path_precision_input, enhance_with_ai_input, remove_bg_input, model_choice_input
         ],
         outputs=[svg_output, download_output]
     )