Update app.py

app.py

@@ -17,53 +17,99 @@ MODEL_OPTIONS = {
 
 loaded_models = {}
 
+def debug_editor_data(editor_data):
+    """Debug function to understand the structure of ImageEditor data"""
+    if editor_data is None:
+        return "No data provided"
+    
+    info = f"Type: {type(editor_data)}\n"
+    
+    if isinstance(editor_data, dict):
+        info += f"Keys: {list(editor_data.keys())}\n"
+        for key, value in editor_data.items():
+            info += f"  {key}: {type(value)}\n"
+            if key == 'layers' and value:
+                info += f"    Layers count: {len(value)}\n"
+                for i, layer in enumerate(value):
+                    info += f"    Layer {i}: {type(layer)}\n"
+                    if hasattr(layer, 'size'):
+                        info += f"      Size: {layer.size}\n"
+                        info += f"      Mode: {layer.mode}\n"
+    else:
+        if hasattr(editor_data, 'size'):
+            info += f"Size: {editor_data.size}\n"
+            info += f"Mode: {editor_data.mode}\n"
+    
+    return info
+
 def extract_region_from_editor(editor_data):
     """Extract the selected region from ImageEditor data"""
     if editor_data is None:
         return None, None
     
-    # Get the background image and layers
-    background = editor_data.get('background')
-    layers = editor_data.get('layers', [])
+    # ImageEditor can return different formats depending on version
+    # Let's handle the most common cases
     
-    if background is None:
-        return None, None
-    
-    # Convert background to numpy array
-    background_np = np.array(background)
-    
-    # If there are layers (selections), process them
-    if layers:
-        # For simplicity, we'll use the first layer as the selection mask
-        # In a more complex implementation, you might want to combine multiple layers
-        selection_layer = layers[0]
+    if isinstance(editor_data, dict):
+        # Case 1: Dictionary with 'background' and 'layers'
+        background = editor_data.get('background')
+        layers = editor_data.get('layers', [])
+        
+        if background is None:
+            return None, None
+            
+        # Convert background to numpy array
+        background_np = np.array(background)
         
-        # The layer contains the selection information
-        # This is a simplified approach - you might need to adjust based on your specific needs
-        selection_image = selection_layer.get('image')
-        if selection_image is not None:
-            # Convert selection to mask
-            selection_np = np.array(selection_image)
+        # If there are drawn layers, try to extract selection
+        if layers and len(layers) > 0:
+            # Get the first layer
+            selection_layer = layers[0]
             
-            # Create a binary mask from the selection
+            # Convert to numpy array
+            selection_np = np.array(selection_layer)
+            
+            # Find non-transparent/non-black pixels as selection
             if len(selection_np.shape) == 3:
-                # If it's an RGB image, convert to grayscale for mask
-                mask = cv2.cvtColor(selection_np, cv2.COLOR_RGB2GRAY)
+                # For RGB, look for non-black pixels
+                if selection_np.shape[2] == 4:  # RGBA
+                    # Use alpha channel
+                    mask = selection_np[:, :, 3] > 0
+                else:  # RGB
+                    # Use non-black pixels
+                    mask = np.any(selection_np > 0, axis=2)
             else:
-                mask = selection_np
+                # Grayscale
+                mask = selection_np > 0
             
             # Find bounding box of the selection
-            coords = np.where(mask > 0)
+            coords = np.where(mask)
             if len(coords[0]) > 0:
                 y_min, y_max = coords[0].min(), coords[0].max()
                 x_min, x_max = coords[1].min(), coords[1].max()
                 
+                # Add some padding to ensure we don't get tiny regions
+                pad = 5
+                h, w = background_np.shape[:2]
+                y_min = max(0, y_min - pad)
+                y_max = min(h, y_max + pad)
+                x_min = max(0, x_min - pad)
+                x_max = min(w, x_max + pad)
+                
                 # Extract the region
                 region = background_np[y_min:y_max+1, x_min:x_max+1]
                 return region, (x_min, y_min, x_max, y_max)
+        
+        # If no selection, return the full image
+        return background_np, None
     
-    # If no selection, return the full image
-    return background_np, None
+    else:
+        # Case 2: Direct PIL Image
+        if hasattr(editor_data, 'size'):  # Check if it's a PIL Image
+            image_np = np.array(editor_data)
+            return image_np, None
+        else:
+            return None, None
 
 def crop_image_with_coords(image_np, coords):
     """Crop image using provided coordinates"""
@@ -83,83 +129,90 @@ def segment_and_count(editor_data, model_choice, crop_coords=None):
         model_choice: Selected model for segmentation
         crop_coords: Optional manual crop coordinates as "x_min,y_min,x_max,y_max"
     """
-    # Load the model
-    model_filename = MODEL_OPTIONS[model_choice]
-    model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
-    
-    if model_filename in loaded_models:
-        model = loaded_models[model_filename]
-    else:
-        model = models.CellposeModel(gpu=True, pretrained_model=model_path)
-        loaded_models[model_filename] = model
-    
-    # Extract region from editor
-    region_np, region_coords = extract_region_from_editor(editor_data)
-    
-    if region_np is None:
-        return 0, None, "No image provided"
-    
-    # If manual crop coordinates are provided, use them instead
-    if crop_coords and crop_coords.strip():
-        try:
-            coords = [int(x.strip()) for x in crop_coords.split(',')]
-            if len(coords) == 4:
-                x_min, y_min, x_max, y_max = coords
-                # Ensure coordinates are within image bounds
-                h, w = region_np.shape[:2]
-                x_min = max(0, min(x_min, w-1))
-                y_min = max(0, min(y_min, h-1))
-                x_max = max(x_min+1, min(x_max, w))
-                y_max = max(y_min+1, min(y_max, h))
-                
-                region_np = region_np[y_min:y_max, x_min:x_max]
-                region_coords = (x_min, y_min, x_max, y_max)
-        except ValueError:
-            pass  # Invalid coordinates, continue with current region
-    
-    # If grayscale, convert to RGB
-    if len(region_np.shape) == 2:
-        region_np = cv2.cvtColor(region_np, cv2.COLOR_GRAY2RGB)
-    elif len(region_np.shape) == 3 and region_np.shape[2] == 4:
-        # Handle RGBA images
-        region_np = cv2.cvtColor(region_np, cv2.COLOR_RGBA2RGB)
-    
-    # Run Cellpose on the selected region
-    masks, flows, styles = model.eval(region_np, diameter=None, channels=[0, 0])
-    
-    # Count unique cells
-    cell_count = len(np.unique(masks)) - 1  # subtract 1 for background
-    
-    # Create better overlay visualization
-    overlay = region_np.copy().astype(np.float32)
-    
-    # Create colored mask overlay
-    if masks.max() > 0:
-        # Generate random colors for each cell
-        np.random.seed(42)  # For reproducible colors
-        colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
-        colors[0] = [0, 0, 0]  # Background stays black
+    try:
+        # Debug info
+        debug_info = debug_editor_data(editor_data)
         
-        # Create colored overlay
-        colored_mask = colors[masks]
+        # Load the model
+        model_filename = MODEL_OPTIONS[model_choice]
+        model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
         
-        # Blend with original image
-        alpha = 0.4
-        overlay = (1 - alpha) * overlay + alpha * colored_mask
-    
-    # Ensure values are in valid range and convert to uint8
-    overlay = np.clip(overlay, 0, 255).astype(np.uint8)
-    
-    # Convert result to PIL Image for output
-    overlay_image = Image.fromarray(overlay)
-    
-    # Create info message
-    if region_coords:
-        info_msg = f"Processed region: {region_coords[0]},{region_coords[1]} to {region_coords[2]},{region_coords[3]}"
-    else:
-        info_msg = "Processed entire image"
-    
-    return cell_count, overlay_image, info_msg
+        if model_filename in loaded_models:
+            model = loaded_models[model_filename]
+        else:
+            model = models.CellposeModel(gpu=True, pretrained_model=model_path)
+            loaded_models[model_filename] = model
+        
+        # Extract region from editor
+        region_np, region_coords = extract_region_from_editor(editor_data)
+        
+        if region_np is None:
+            return 0, None, f"No image provided. Debug info:\n{debug_info}"
+        
+        # If manual crop coordinates are provided, use them instead
+        if crop_coords and crop_coords.strip():
+            try:
+                coords = [int(x.strip()) for x in crop_coords.split(',')]
+                if len(coords) == 4:
+                    x_min, y_min, x_max, y_max = coords
+                    # Ensure coordinates are within image bounds
+                    h, w = region_np.shape[:2]
+                    x_min = max(0, min(x_min, w-1))
+                    y_min = max(0, min(y_min, h-1))
+                    x_max = max(x_min+1, min(x_max, w))
+                    y_max = max(y_min+1, min(y_max, h))
+                    
+                    region_np = region_np[y_min:y_max, x_min:x_max]
+                    region_coords = (x_min, y_min, x_max, y_max)
+            except ValueError:
+                pass  # Invalid coordinates, continue with current region
+        
+        # If grayscale, convert to RGB
+        if len(region_np.shape) == 2:
+            region_np = cv2.cvtColor(region_np, cv2.COLOR_GRAY2RGB)
+        elif len(region_np.shape) == 3 and region_np.shape[2] == 4:
+            # Handle RGBA images
+            region_np = cv2.cvtColor(region_np, cv2.COLOR_RGBA2RGB)
+        
+        # Run Cellpose on the selected region
+        masks, flows, styles = model.eval(region_np, diameter=None, channels=[0, 0])
+        
+        # Count unique cells
+        cell_count = len(np.unique(masks)) - 1  # subtract 1 for background
+        
+        # Create better overlay visualization
+        overlay = region_np.copy().astype(np.float32)
+        
+        # Create colored mask overlay
+        if masks.max() > 0:
+            # Generate random colors for each cell
+            np.random.seed(42)  # For reproducible colors
+            colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
+            colors[0] = [0, 0, 0]  # Background stays black
+            
+            # Create colored overlay
+            colored_mask = colors[masks]
+            
+            # Blend with original image
+            alpha = 0.4
+            overlay = (1 - alpha) * overlay + alpha * colored_mask
+        
+        # Ensure values are in valid range and convert to uint8
+        overlay = np.clip(overlay, 0, 255).astype(np.uint8)
+        
+        # Convert result to PIL Image for output
+        overlay_image = Image.fromarray(overlay)
+        
+        # Create info message
+        if region_coords:
+            info_msg = f"Processed region: {region_coords[0]},{region_coords[1]} to {region_coords[2]},{region_coords[3]}\nRegion size: {region_np.shape}\nDebug: {debug_info}"
+        else:
+            info_msg = f"Processed entire image\nImage size: {region_np.shape}\nDebug: {debug_info}"
+        
+        return cell_count, overlay_image, info_msg
+        
+    except Exception as e:
+        return 0, None, f"Error occurred: {str(e)}\nDebug info:\n{debug_info if 'debug_info' in locals() else 'No debug info'}"
 
 # Alternative function for simple image input with coordinate textbox
 @spaces.GPU
@@ -293,4 +346,3 @@ with gr.Blocks(title="Cell Counter with Region Selection") as demo:
 if __name__ == "__main__":
     demo.launch()
 
-