Update app.py

app.py

@@ -24,7 +24,17 @@ def get_sam_masks(image):
     else:
         image_np = image
         image_pil = Image.fromarray(image_np)
-        
+    
+    # Handle grayscale images
+    if len(image_np.shape) == 2:
+        # Convert grayscale to RGB
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
+        image_pil = Image.fromarray(image_np)
+    elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
+        # Convert RGBA to RGB
+        image_np = image_np[:, :, :3]
+        image_pil = Image.fromarray(image_np)
+    
     h, w = image_np.shape[:2]
     
     # Create a grid of points to sample the image
@@ -52,20 +62,21 @@ def get_sam_masks(image):
             inputs["reshaped_input_sizes"].cpu()
         )
     
-    # Combine all masks to create importance map
+    # Combine all masks to create importance map (ensure 2D)
     importance_map = np.zeros((h, w), dtype=np.float32)
-    individual_masks = []
     
     for i in range(len(masks[0])):
         mask = masks[0][i].numpy().astype(np.float32)
-        individual_masks.append(mask)
+        # Ensure mask is 2D
+        if len(mask.shape) > 2:
+            mask = mask[:, :, 0]
         importance_map += mask
     
     # Normalize to 0-1
     if importance_map.max() > 0:
         importance_map = importance_map / importance_map.max()
     
-    return importance_map, individual_masks
+    return importance_map
 
 def find_optimal_crop(image, target_ratio, importance_map):
     """Find the optimal crop area that preserves important content while matching target ratio"""
@@ -132,24 +143,23 @@ def apply_crop(image, crop_box):
     else:
         return image.crop((left, top, right, bottom))
 
-def visualize_importance_map(image, importance_map):
+def visualize_importance_map(image_np, importance_map):
     """Create a visualization of the importance map overlaid on the image"""
-    # Make sure importance_map is 2D
+    # Ensure importance_map is 2D
     if len(importance_map.shape) > 2:
-        importance_map = importance_map[:, :, 0]  # Take first channel if multichannel
+        importance_map = importance_map[:, :, 0]
         
-    # Resize importance map to match image dimensions if needed
-    if importance_map.shape[:2] != image.shape[:2]:
-        importance_map = cv2.resize(importance_map, (image.shape[1], image.shape[0]))
+    # Ensure image is RGB
+    if len(image_np.shape) == 2:
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
+    elif image_np.shape[2] == 4:
+        image_np = image_np[:, :, :3]
     
     # Convert to colormap
     heatmap = cv2.applyColorMap((importance_map * 255).astype(np.uint8), cv2.COLORMAP_JET)
     
-    # Convert image to BGR if it's RGB
-    if len(image.shape) == 3 and image.shape[2] == 3:
-        image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    else:
-        image_bgr = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+    # Convert image to BGR for OpenCV
+    image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
     
     # Overlay heatmap on original image
     overlay = cv2.addWeighted(image_bgr, 0.7, heatmap, 0.3, 0)
@@ -161,21 +171,22 @@ def visualize_importance_map(image, importance_map):
 
 def adjust_aspect_ratio(image, target_ratio):
     """Main function to adjust aspect ratio through intelligent cropping"""
+    # Convert image to numpy if needed
+    if isinstance(image, Image.Image):
+        image_np = np.array(image)
+    else:
+        image_np = image
+    
     # Get segmentation masks and importance map
-    importance_map, _ = get_sam_masks(image)
+    importance_map = get_sam_masks(image_np)
     
     # Find optimal crop box
-    crop_box = find_optimal_crop(image, target_ratio, importance_map)
+    crop_box = find_optimal_crop(image_np, target_ratio, importance_map)
     
     # Apply the crop
-    result = apply_crop(image, crop_box)
+    result = apply_crop(image_np, crop_box)
     
     # Create visualization of importance map
-    if isinstance(image, np.ndarray):
-        image_np = image
-    else:
-        image_np = np.array(image)
-        
     importance_vis = visualize_importance_map(image_np, importance_map)
     
     return result, importance_vis
@@ -251,4 +262,4 @@ with gr.Blocks(title="SAM-Based Smart Crop Aspect Ratio Adjuster") as demo:
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)  # Set share=True to create a public link