Fix Gradio compatibility issues and canvas processing bugs

app.py

@@ -21,14 +21,37 @@ def create_interface():
                 # Draw Region Column
                 with gr.Column():
                     gr.Markdown("""<p style="text-align: center; font-size: 20px">1. Draw Regions</p>""")
-                    canvas = gr.Image(type="numpy", tool="sketch", label=" ", height=CANVAS_SIZE, width=CANVAS_SIZE)
+                    # Use ImageEditor for newer Gradio versions, fallback to Image with brush
+                    try:
+                        canvas = gr.ImageEditor(
+                            label=" ", 
+                            height=CANVAS_SIZE, 
+                            width=CANVAS_SIZE,
+                            brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed")
+                        )
+                    except:
+                        # Fallback for older Gradio versions
+                        canvas = gr.Image(
+                            type="numpy", 
+                            label=" ", 
+                            height=CANVAS_SIZE, 
+                            width=CANVAS_SIZE,
+                            sources=["upload", "webcam", "clipboard"]
+                        )
+                    
                     with gr.Row():
                         fit_btn = gr.Button("Resize Image")
 
                 # Control Points Column
                 with gr.Column():
                     gr.Markdown("""<p style="text-align: center; font-size: 20px">2. Control Points</p>""")
-                    input_img = gr.Image(type="numpy", label=" ", height=CANVAS_SIZE, width=CANVAS_SIZE, interactive=True)
+                    input_img = gr.Image(
+                        type="numpy", 
+                        label=" ", 
+                        height=CANVAS_SIZE, 
+                        width=CANVAS_SIZE, 
+                        interactive=True
+                    )
                     with gr.Row():
                         undo_btn = gr.Button("Undo Point")
                         clear_btn = gr.Button("Clear Points")
@@ -36,14 +59,27 @@ def create_interface():
                 # Results Column
                 with gr.Column():
                     gr.Markdown("""<p style="text-align: center; font-size: 20px">Results</p>""")
-                    output_img = gr.Image(type="numpy", label=" ", height=CANVAS_SIZE, width=CANVAS_SIZE, interactive=False)
+                    output_img = gr.Image(
+                        type="numpy", 
+                        label=" ", 
+                        height=CANVAS_SIZE, 
+                        width=CANVAS_SIZE, 
+                        interactive=False
+                    )
                     with gr.Row():
                         run_btn = gr.Button("Inpaint")
                         reset_btn = gr.Button("Reset All")
 
         # Generation Parameters
         with gr.Row():
-            inpaint_ks = gr.Slider(minimum=0, maximum=25, value=5, step=1, label='How much to expand inpainting mask', interactive=True)
+            inpaint_ks = gr.Slider(
+                minimum=0, 
+                maximum=25, 
+                value=5, 
+                step=1, 
+                label='How much to expand inpainting mask', 
+                interactive=True
+            )
 
         setup_events(
             components={
@@ -95,7 +131,10 @@ def setup_events(components, state, buttons):
 
     # Canvas interaction events
     def setup_canvas_events():
-        components['canvas'].edit(
+        # Handle both ImageEditor and Image events
+        canvas_event = components['canvas'].change if hasattr(components['canvas'], 'change') else components['canvas'].edit
+        
+        canvas_event(
             visualize_user_drag,
             [components['canvas'], state['points_list']],
             [components['input_img']]

utils/ui_utils.py

@@ -25,14 +25,18 @@ pipe = None
 # UI functions
 def clear_all(length):
     """Reset UI by clearing all input images and parameters."""
-    return (gr.Image(value=None, height=length, width=length),) * 3 + ([], 2, None)
+    return (gr.Image(value=None, height=length, width=length),) * 3 + ([], 5, None)
 
 def resize(canvas, gen_length, canvas_length):
     """Resize canvas while maintaining aspect ratio."""
     if not canvas:
         return (gr.Image(value=None, width=canvas_length, height=canvas_length),) * 3
 
-    image = process_canvas(canvas)[0]
+    result = process_canvas(canvas)
+    if result[0] is None:  # Check if image is None
+        return (gr.Image(value=None, width=canvas_length, height=canvas_length),) * 3
+    
+    image = result[0]
     aspect_ratio = image.shape[1] / image.shape[0]
     is_landscape = aspect_ratio >= 1
 
@@ -49,8 +53,61 @@ def resize(canvas, gen_length, canvas_length):
 
 def process_canvas(canvas):
     """Extracts the image (H, W, 3) and the mask (H, W) from a Gradio canvas object."""
-    image = canvas["image"].copy()
-    mask = np.uint8(canvas["mask"][:, :, 0] > 0).copy()
+    # Handle None canvas
+    if canvas is None:
+        return None, None
+    
+    # Handle new ImageEditor format
+    if isinstance(canvas, dict):
+        if 'background' in canvas and 'layers' in canvas:
+            # New ImageEditor format
+            if canvas["background"] is None:
+                return None, None
+            image = canvas["background"].copy()
+            
+            # Ensure image is 3-channel RGB
+            if len(image.shape) == 3 and image.shape[2] == 4:
+                image = image[:, :, :3]  # Remove alpha channel
+            elif len(image.shape) == 2:
+                image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+            
+            # Try to extract mask from layers
+            mask = np.zeros(image.shape[:2], dtype=np.uint8)
+            if canvas["layers"]:
+                for layer in canvas["layers"]:
+                    if isinstance(layer, np.ndarray) and len(layer.shape) >= 2:
+                        layer_mask = np.uint8(layer[:, :, 0] > 0) if len(layer.shape) == 3 else np.uint8(layer > 0)
+                        mask = np.logical_or(mask, layer_mask).astype(np.uint8)
+        elif 'image' in canvas and 'mask' in canvas:
+            # Old format
+            if canvas["image"] is None:
+                return None, None
+            image = canvas["image"].copy()
+            
+            # Ensure image is 3-channel RGB
+            if len(image.shape) == 3 and image.shape[2] == 4:
+                image = image[:, :, :3]  # Remove alpha channel
+            elif len(image.shape) == 2:
+                image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+            
+            mask = np.uint8(canvas["mask"][:, :, 0] > 0).copy() if canvas["mask"] is not None else np.zeros(image.shape[:2], dtype=np.uint8)
+        else:
+            # Fallback
+            return None, None
+    else:
+        # Direct numpy array
+        if canvas is None:
+            return None, None
+        image = canvas.copy() if isinstance(canvas, np.ndarray) else np.array(canvas)
+        
+        # Ensure image is 3-channel RGB
+        if len(image.shape) == 3 and image.shape[2] == 4:
+            image = image[:, :, :3]  # Remove alpha channel
+        elif len(image.shape) == 2:
+            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+        
+        mask = np.zeros(image.shape[:2], dtype=np.uint8)
+    
     return image, mask
 
 # Point manipulation functions
@@ -82,22 +139,37 @@ def visualize_user_drag(canvas, points):
     if canvas is None:
         return None
     
-    image, mask = process_canvas(canvas)
+    result = process_canvas(canvas)
+    if result[0] is None:  # Check if image is None
+        return None
+    
+    image, mask = result
+
+    # Ensure image is uint8 and 3-channel
+    if image.dtype != np.uint8:
+        image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8)
+    
+    if len(image.shape) != 3 or image.shape[2] != 3:
+        return None
 
     # Apply colored mask overlay
-    result = image.copy()
-    result[mask == 1] = [255, 0, 0]  # Red color
-    image = cv2.addWeighted(result, 0.3, image, 0.7, 0)
+    result_img = image.copy()
+    if np.any(mask == 1):
+        result_img[mask == 1] = [255, 0, 0]  # Red color
+        image = cv2.addWeighted(result_img, 0.3, image, 0.7, 0)
     
     # Draw mask outline
-    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    cv2.drawContours(image, contours, -1, (255, 255, 255), 2)
+    if np.any(mask > 0):
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        cv2.drawContours(image, contours, -1, (255, 255, 255), 2)
     
     # Draw control points and motion vectors
+    prev_point = None
     for idx, point in enumerate(points, 1):
         if idx % 2 == 0:
             cv2.circle(image, tuple(point), 10, (0, 0, 255), -1)  # End point
-            cv2.arrowedLine(image, prev_point, point, (255, 255, 255), 4, tipLength=0.5)
+            if prev_point is not None:
+                cv2.arrowedLine(image, prev_point, point, (255, 255, 255), 4, tipLength=0.5)
         else:
             cv2.circle(image, tuple(point), 10, (255, 0, 0), -1)  # Start point
             prev_point = point
@@ -109,7 +181,19 @@ def preview_out_image(canvas, points, inpaint_ks):
     if canvas is None:
         return None, None
     
-    image, mask = process_canvas(canvas)
+    result = process_canvas(canvas)
+    if result[0] is None:  # Check if image is None
+        return None, None
+    
+    image, mask = result
+    
+    # Ensure image is uint8 and 3-channel
+    if image.dtype != np.uint8:
+        image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8)
+    
+    if len(image.shape) != 3 or image.shape[2] != 3:
+        return image, None
+    
     if len(points) < 2:
         return image, None
     
@@ -120,15 +204,19 @@ def preview_out_image(canvas, points, inpaint_ks):
         gr.Warning('Click Resize Image Button first.')
         return image, None
     
-    handle_pts, target_pts, inpaint_mask = bi_warp(mask, points, inpaint_ks)
-    image[target_pts[:, 1], target_pts[:, 0]] = image[handle_pts[:, 1], handle_pts[:, 0]]
+    try:
+        handle_pts, target_pts, inpaint_mask = bi_warp(mask, points, inpaint_ks)
+        image[target_pts[:, 1], target_pts[:, 0]] = image[handle_pts[:, 1], handle_pts[:, 0]]
 
-    # Add grid pattern to highlight inpainting regions
-    background = np.ones_like(mask) * 255
-    background[::10] = background[:, ::10] = 0
-    image = np.where(inpaint_mask[..., np.newaxis]==1, background[..., np.newaxis], image)
-    
-    return image, (inpaint_mask * 255).astype(np.uint8)
+        # Add grid pattern to highlight inpainting regions
+        background = np.ones_like(mask) * 255
+        background[::10] = background[:, ::10] = 0
+        image = np.where(inpaint_mask[..., np.newaxis]==1, background[..., np.newaxis], image)
+        
+        return image, (inpaint_mask * 255).astype(np.uint8)
+    except Exception as e:
+        gr.Warning(f"Preview failed: {str(e)}")
+        return image, None
 
 # Inpaint tools
 def setup_pipeline(device='cuda', model_version='v1-5'):