Update app.py

app.py

@@ -132,6 +132,33 @@ def apply_crop(image, crop_box):
     else:
         return image.crop((left, top, right, bottom))
 
+def visualize_importance_map(image, importance_map):
+    """Create a visualization of the importance map overlaid on the image"""
+    # Make sure importance_map is 2D
+    if len(importance_map.shape) > 2:
+        importance_map = importance_map[:, :, 0]  # Take first channel if multichannel
+        
+    # 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]))
+    
+    # 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)
+    
+    # Overlay heatmap on original image
+    overlay = cv2.addWeighted(image_bgr, 0.7, heatmap, 0.3, 0)
+    
+    # Convert back to RGB for display
+    overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
+    
+    return overlay_rgb
+
 def adjust_aspect_ratio(image, target_ratio):
     """Main function to adjust aspect ratio through intelligent cropping"""
     # Get segmentation masks and importance map
@@ -143,7 +170,15 @@ def adjust_aspect_ratio(image, target_ratio):
     # Apply the crop
     result = apply_crop(image, crop_box)
     
-    return result, importance_map
+    # 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
 
 def process_image(input_image, target_ratio="16:9"):
     """Process function for Gradio interface"""
@@ -155,7 +190,7 @@ def process_image(input_image, target_ratio="16:9"):
             image = input_image
             
         # Adjust aspect ratio
-        result, importance_map = adjust_aspect_ratio(image, target_ratio)
+        result, importance_vis = adjust_aspect_ratio(image, target_ratio)
         
         # Convert result to appropriate format
         if isinstance(result, np.ndarray):
@@ -163,21 +198,15 @@ def process_image(input_image, target_ratio="16:9"):
         else:
             result_pil = result
         
-        # Visualize importance map for debugging
-        if isinstance(importance_map, np.ndarray):
-            # Convert to heatmap
-            heatmap = (importance_map * 255).astype(np.uint8)
-            heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-            
-            # Convert to PIL
-            heatmap_pil = Image.fromarray(heatmap)
-            
-            return [result_pil, heatmap_pil]
+        # Convert importance visualization to PIL
+        importance_vis_pil = Image.fromarray(importance_vis)
         
-        return [result_pil, None]
+        return [result_pil, importance_vis_pil]
     
     except Exception as e:
-        print(f"Error processing image: {e}")
+        print(f"Error processing image: {str(e)}")
+        import traceback
+        traceback.print_exc()
         return [None, None]
 
 # Create the Gradio interface
@@ -200,7 +229,7 @@ with gr.Blocks(title="SAM-Based Smart Crop Aspect Ratio Adjuster") as demo:
         
         with gr.Column():
             output_image = gr.Image(label="Processed Image")
-            importance_map_vis = gr.Image(label="Importance Map (Debug View)")
+            importance_map_vis = gr.Image(label="Importance Map (Overlay View)")
     
     submit_btn.click(
         process_image, 
@@ -211,7 +240,7 @@ with gr.Blocks(title="SAM-Based Smart Crop Aspect Ratio Adjuster") as demo:
     gr.Markdown("""
     ## How it works
     1. **Segmentation**: Uses Meta's Segment Anything Model (SAM) to identify important regions in your image
-    2. **Importance Mapping**: Creates a heatmap of important areas based on segmentation masks
+    2. **Importance Mapping**: Creates a heatmap of important areas based on segmentation masks (shown in the overlay)
     3. **Smart Cropping**: Finds the optimal crop window that preserves the most important content
     
     ## Tips