Update app.py

app.py

@@ -1,71 +1,74 @@
-import io
 import gradio as gr
-from PIL import Image
-import vtracer
-import tempfile
+import torch
+from torchvision import transforms
+from PIL import Image, ImageEnhance
+import rembg
 
-def convert_image(image, 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."""
-
-    # 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()
-
-    return gr.HTML(f'<svg viewBox="0 0 {image.width} {image.height}">{svg_str}</svg>'), temp_file.name
+# Upscaling function using a basic upscaling approach
+def upscale_image(image, upscale_factor=2, sharpness=1.0, contrast=1.0, brightness=1.0):
+    # Resize the image
+    width, height = image.size
+    new_size = (int(width * upscale_factor), int(height * upscale_factor))
+    upscaled_image = image.resize(new_size, Image.BICUBIC)
+    
+    # Apply sharpness, contrast, and brightness adjustments
+    upscaled_image = ImageEnhance.Sharpness(upscaled_image).enhance(sharpness)
+    upscaled_image = ImageEnhance.Contrast(upscaled_image).enhance(contrast)
+    upscaled_image = ImageEnhance.Brightness(upscaled_image).enhance(brightness)
+    
+    return upscaled_image
 
-# Gradio interface
-iface = gr.Blocks()
+# Remastering function with background removal and enhancements
+def remaster_image(image, remove_background=False, color_balance=1.0, detail_enhancement=1.0):
+    if remove_background:
+        image = rembg.remove(image)
+    
+    # Apply color balance and detail enhancement
+    image = ImageEnhance.Color(image).enhance(color_balance)
+    image = ImageEnhance.Sharpness(image).enhance(detail_enhancement)
+    
+    return image
 
-with iface:
+# Main function combining upscaling and remastering
+def process_image(image, apply_upscale=False, upscale_factor=2, sharpness=1.0, contrast=1.0, brightness=1.0,
+                  apply_remaster=False, remove_background=False, color_balance=1.0, detail_enhancement=1.0):
+    if apply_upscale:
+        image = upscale_image(image, upscale_factor, sharpness, contrast, brightness)
     
-    gr.Interface(
-    fn=convert_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Radio(choices=["Color", "Binary"], value="Color", label="Color Mode"),
-        gr.Radio(choices=["Stacked", "Cutout"], value="Stacked", label="Hierarchical"),
-        gr.Radio(choices=["Spline", "Polygon", "None"], value="Spline", label="Mode"),
-        gr.Slider(minimum=1, maximum=100, value=4, step=1, label="Filter Speckle"),
-        gr.Slider(minimum=1, maximum=100, value=6, step=1, label="Color Precision"),
-        gr.Slider(minimum=1, maximum=100, value=16, step=1, label="Layer Difference"),
-        gr.Slider(minimum=1, maximum=100, value=60, step=1, label="Corner Threshold"),
-        gr.Slider(minimum=1, maximum=100, value=4.0, step=0.5, label="Length Threshold"),
-        gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Max Iterations"),
-        gr.Slider(minimum=1, maximum=100, value=45, step=1, label="Splice Threshold"),
-        gr.Slider(minimum=1, maximum=100, value=8, step=1, label="Path Precision")
-    ],
-    outputs=[
-        gr.HTML(label="SVG Output"),
-        gr.File(label="Download SVG")
-    ],
-    title="Convert Image to SVG vectors",
-    description="Upload an image and customize the conversion parameters as needed.",
-)
+    if apply_remaster:
+        image = remaster_image(image, remove_background, color_balance, detail_enhancement)
+    
+    return image
 
+# Gradio UI
+with gr.Blocks() as demo:
+    with gr.Row():
+        image_input = gr.Image(label="Upload Image", type="pil")
+        image_output = gr.Image(label="Output Image")
+    
+    with gr.Row():
+        with gr.Group():
+            gr.Markdown("### Upscaling Options")
+            upscale_checkbox = gr.Checkbox(label="Apply Upscaling")
+            upscale_factor = gr.Slider(1, 4, value=2, label="Upscale Factor")
+            sharpness = gr.Slider(0.5, 2.0, value=1.0, label="Sharpness")
+            contrast = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
+            brightness = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
+        
+        with gr.Group():
+            gr.Markdown("### Remastering Options")
+            remaster_checkbox = gr.Checkbox(label="Apply Remastering")
+            remove_background = gr.Checkbox(label="Remove Background")
+            color_balance = gr.Slider(0.5, 2.0, value=1.0, label="Color Balance")
+            detail_enhancement = gr.Slider(0.5, 2.0, value=1.0, label="Detail Enhancement")
     
-iface.launch()
+    process_button = gr.Button("Process Image")
+    
+    process_button.click(
+        process_image,
+        inputs=[image_input, upscale_checkbox, upscale_factor, sharpness, contrast, brightness,
+                remaster_checkbox, remove_background, color_balance, detail_enhancement],
+        outputs=image_output
+    )
+
+demo.launch()