app.py

import gradio as gr
from PIL import Image, ImageFilter, ImageOps
import numpy as np
import io
import tempfile
import vtracer
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, 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)

    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 with rembg for smoothing edges and improving vectorization
        image_np = remove(image_np)
        image = Image.fromarray(image_np)
    
    return 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,
                  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,
        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

# 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]
    )

iface.launch()