Upload app.py with huggingface_hub

app.py

@@ -7,9 +7,13 @@ Original file is located at
     https://colab.research.google.com/drive/1GoN403JoCKCLavITeCYoMzFXKctQ83en
 """
 
+from IPython import get_ipython
+from IPython.display import display
+# %%
 import gradio as gr
 from PIL import Image
 import numpy as np
+import cv2  # Import OpenCV
 
 def sanitize_hex_color(color):
     """
@@ -35,7 +39,8 @@ def sanitize_hex_color(color):
 
 def remove_background_color(image, hex_color, tolerance):
     """
-    Remove pixels of a specified hex color from an image.
+    Remove pixels of a specified hex color from an image,
+    with edge detection for refinement.
     """
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
@@ -47,25 +52,20 @@ def remove_background_color(image, hex_color, tolerance):
         raise ValueError(str(e))
     data = np.array(image)
     mask = np.all(np.abs(data[:,:,:3] - target_color) <= tolerance, axis=2)
+
+    # Edge detection using Canny
+    edges = cv2.Canny(cv2.cvtColor(data, cv2.COLOR_RGBA2GRAY), 100, 200)
+    # Adjust thresholds (100, 200) as needed
+
+    # Refine mask based on edges
+    kernel = np.ones((3, 3), np.uint8)
+    dilated_edges = cv2.dilate(edges, kernel, iterations=1)
+    mask = np.logical_or(mask, dilated_edges)
+
     alpha_channel = np.where(mask, 0, 255)
     data[:, :, 3] = alpha_channel
     return data
 
-
-def color_remover_interface(image, hex_color, tolerance, _=None):
-    """
-    Gradio interface function for color removal.
-    """
-    try:
-        if image is None:
-            raise ValueError("Please upload an image")
-        hex_color = sanitize_hex_color(hex_color)
-        processed_image = remove_background_color(image, hex_color, tolerance)
-        return processed_image  # Return only the processed image
-    except Exception as e:
-        gr.Warning(str(e))
-        return image  # Return original image in case of error
-
 # Deploy using Gradio Client
 iface = gr.Interface(
     fn=color_remover_interface,