Update app.py

app.py

@@ -1,4 +1,5 @@
 # app.py
+
 import gradio as gr
 import numpy as np
 import cv2
@@ -7,24 +8,20 @@ import json
 import time
 from ultralytics import YOLO
 import torch
-from typing import List, Dict, Tuple, Optional
-import base64
-import io
 import os
 
-# Do NOT try to write to ~/.config/Ultralytics on HF Spaces
-os.environ["YOLO_CONFIG_DIR"] = "/tmp"
+# Fix Ultralytics config path for Hugging Face Spaces (read-only /home/user)
+os.environ['YOLO_CONFIG_DIR'] = '/tmp'
 
-# Global variables
+# Global variables to maintain state
 model = None
 
-# Selection Manager class
 class SelectionManager:
     def __init__(self):
         self.selections = []
         self.next_id = 0
         self.active_id = None
-
+    
     def add_selection(self, x, y, width, height, selection_type="manual", label="", confidence=None):
         selection = {
             'id': self.next_id,
@@ -40,19 +37,19 @@ class SelectionManager:
         self.selections.append(selection)
         self.next_id += 1
         return selection['id']
-
+    
     def remove_selection(self, selection_id):
         self.selections = [s for s in self.selections if s['id'] != selection_id]
         if self.active_id == selection_id:
             self.active_id = None
-
+    
     def clear_all(self):
         self.selections = []
         self.active_id = None
-
+    
     def clear_type(self, selection_type):
         self.selections = [s for s in self.selections if s['type'] != selection_type]
-
+    
     def get_total_area(self, img_width, img_height):
         total_selected_area = sum(s['width'] * s['height'] for s in self.selections)
         total_image_area = img_width * img_height
@@ -61,34 +58,32 @@ class SelectionManager:
 # Initialize selection manager
 selection_manager = SelectionManager()
 
-# Load YOLO model
 def load_yolo_model():
     global model
     try:
-        model = YOLO('yolov8n.pt')
+        model = YOLO('yolov8n.pt')  # This will download the model if not cached
         return "✅ YOLO Model loaded successfully"
     except Exception as e:
         return f"❌ Error loading YOLO model: {str(e)}"
 
-# Detect objects
 def detect_objects(image, confidence_threshold=0.5, merge_with_existing=True):
     global model, selection_manager
-
     if model is None:
         return image, "❌ Model not loaded", get_analysis_text(image)
-
+    
     if image is None:
         return None, "❌ No image provided", ""
-
+    
     try:
         img_array = np.array(image)
+        
         start_time = time.time()
         results = model(img_array, conf=confidence_threshold, verbose=False)
-        detection_time = (time.time() - start_time) * 1000
-
+        detection_time = (time.time() - start_time) * 1000  # ms
+        
         if not merge_with_existing:
             selection_manager.clear_type('yolo')
-
+        
         detections_added = 0
         for result in results:
             boxes = result.boxes
@@ -98,7 +93,7 @@ def detect_objects(image, confidence_threshold=0.5, merge_with_existing=True):
                     confidence = box.conf[0].cpu().numpy()
                     class_id = int(box.cls[0].cpu().numpy())
                     class_name = model.names[class_id]
-
+                    
                     selection_manager.add_selection(
                         x=int(x1),
                         y=int(y1),
@@ -109,57 +104,60 @@ def detect_objects(image, confidence_threshold=0.5, merge_with_existing=True):
                         confidence=confidence
                     )
                     detections_added += 1
-
+        
         annotated_image = draw_selections(image)
         status_msg = f"✅ Detected {detections_added} objects in {detection_time:.1f}ms"
         analysis_text = get_analysis_text(image)
         return annotated_image, status_msg, analysis_text
-
+    
     except Exception as e:
         return image, f"❌ Detection error: {str(e)}", get_analysis_text(image)
 
-# Draw selections
 def draw_selections(image):
     if image is None:
         return None
-
+    
     img_copy = image.copy()
     draw = ImageDraw.Draw(img_copy)
-
+    
     try:
         font = ImageFont.truetype("arial.ttf", 12)
     except:
         font = ImageFont.load_default()
-
+    
     for selection in selection_manager.selections:
         x, y, w, h = selection['x'], selection['y'], selection['width'], selection['height']
-        outline_color = 'green' if selection['type'] == 'yolo' else 'blue'
-        fill_color = (0, 255, 0, 60) if selection['type'] == 'yolo' else (0, 0, 255, 60)
-
+        
+        if selection['type'] == 'yolo':
+            outline_color = 'green'
+            fill_color = (0, 255, 0, 60)
+        else:
+            outline_color = 'blue'
+            fill_color = (0, 0, 255, 60)
+        
         draw.rectangle([x, y, x + w, y + h], outline=outline_color, width=2)
-
+        
         overlay = Image.new('RGBA', img_copy.size, (0, 0, 0, 0))
         overlay_draw = ImageDraw.Draw(overlay)
         overlay_draw.rectangle([x, y, x + w, y + h], fill=fill_color)
         img_copy = Image.alpha_composite(img_copy.convert('RGBA'), overlay).convert('RGB')
-
+        
         if selection['label']:
             label_y = y - 15 if y > 15 else y + h + 5
             draw.text((x, label_y), selection['label'], fill=outline_color, font=font)
-
+    
     return img_copy
 
-# Analysis text
 def get_analysis_text(image):
     if image is None:
         return "No image loaded"
-
+    
     img_width, img_height = image.size
     total_selected_area, total_image_area = selection_manager.get_total_area(img_width, img_height)
-
+    
     selected_percentage = (total_selected_area / total_image_area) * 100 if total_image_area > 0 else 0
     extra_percentage = 100 - selected_percentage
-
+    
     analysis = f"""
 ## 📊 Space Analysis
 
@@ -167,11 +165,10 @@ def get_analysis_text(image):
 **Total Selected Area:** {total_selected_area:,} px²  
 **Selected Region:** {selected_percentage:.1f}% of total  
 **Extra Space:** {extra_percentage:.1f}% of total  
-**Number of Selections:** {len(selection_manager.selections)}
+**Number of Selections:** {len(selection_manager.selections)}  
 
 ### 📋 Selection Details:
 """
-
     for i, selection in enumerate(selection_manager.selections, 1):
         area = selection['width'] * selection['height']
         analysis += f"""
@@ -180,35 +177,35 @@ def get_analysis_text(image):
 - Area: {area:,} px²  
 - Dimensions: {selection['width']}×{selection['height']}  
 """
-        if selection['confidence']:
+        if selection['confidence'] is not None:
             analysis += f"- Confidence: {selection['confidence']:.2f}\n"
-
+    
     return analysis
 
-# Manual selection
 def add_manual_selection(image, selection_data):
     if image is None:
         return image, "❌ No image loaded", ""
-
+    
     try:
         coords = [int(x.strip()) for x in selection_data.split(',')]
         if len(coords) != 4:
             raise ValueError("Invalid format")
-
+        
         x, y, width, height = coords
+        
         img_width, img_height = image.size
         if x < 0 or y < 0 or x + width > img_width or y + height > img_height:
             return image, "❌ Selection coordinates out of bounds", get_analysis_text(image)
-
-        sel_id = selection_manager.add_selection(x, y, width, height, "manual", f"Manual {selection_manager.next_id}")
+        
+        sel_id = selection_manager.add_selection(x, y, width, height, "manual", f"Manual Selection {selection_manager.next_id}")
+        
         annotated_image = draw_selections(image)
         analysis_text = get_analysis_text(image)
         return annotated_image, "✅ Manual selection added", analysis_text
-
+    
     except Exception as e:
         return image, f"❌ Error adding selection: {str(e)}", get_analysis_text(image)
 
-# Clear selections
 def clear_all_selections(image):
     selection_manager.clear_all()
     if image is not None:
@@ -229,41 +226,54 @@ def clear_manual_selections(image):
         return annotated_image, "✅ Manual selections cleared", get_analysis_text(image)
     return None, "✅ Manual selections cleared", ""
 
-# Process image upload
 def process_image_upload(image):
     if image is None:
         return None, "❌ No image uploaded", ""
-
+    
     selection_manager.clear_all()
     analysis_text = get_analysis_text(image)
     return image, "✅ Image loaded successfully", analysis_text
 
-# Gradio Interface
-with gr.Blocks(title="Image Space Analyzer with YOLO") as demo:
-    gr.Markdown("# 🔍 Image Space Analyzer with YOLO")
-    with gr.Row():
-        with gr.Column(scale=2):
-            image_input = gr.Image(type="pil", label="📸 Upload Image", height=500)
-            status_output = gr.Textbox(label="📋 Status", interactive=False, max_lines=2)
-        with gr.Column(scale=1):
-            model_status = gr.Textbox(label="🤖 Model Status", value="Loading YOLO model...", interactive=False)
-            confidence_slider = gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Confidence Threshold")
-            merge_checkbox = gr.Checkbox(label="Merge with existing selections", value=True)
-            detect_btn = gr.Button("🔍 Detect Objects")
-            manual_input = gr.Textbox(label="Selection (x,y,width,height)", placeholder="100,100,200,150")
-            add_manual_btn = gr.Button("➕ Add Manual Selection")
-            clear_all_btn = gr.Button("🗑️ Clear All")
-            clear_yolo_btn = gr.Button("🗑️ Clear YOLO")
-            clear_manual_btn = gr.Button("🗑️ Clear Manual")
-
-    analysis_output = gr.Markdown(label="📊 Analysis Results")
-
-    image_input.upload(process_image_upload, inputs=[image_input], outputs=[image_input, status_output, analysis_output])
-    detect_btn.click(detect_objects, inputs=[image_input, confidence_slider, merge_checkbox], outputs=[image_input, status_output, analysis_output])
-    add_manual_btn.click(add_manual_selection, inputs=[image_input, manual_input], outputs=[image_input, status_output, analysis_output])
-    clear_all_btn.click(clear_all_selections, inputs=[image_input], outputs=[image_input, status_output, analysis_output])
-    clear_yolo_btn.click(clear_yolo_selections, inputs=[image_input], outputs=[image_input, status_output, analysis_output])
-    clear_manual_btn.click(clear_manual_selections, inputs=[image_input], outputs=[image_input], outputs=[image_input, status_output, analysis_output])
-    demo.load(load_yolo_model, outputs=[model_status])
-
-# If you want to deploy in HF Spaces: Do NOT add `launch()` block! HF Spaces will auto launch this file.
+# Gradio interface
+def create_interface():
+    with gr.Blocks(title="Image Space Analyzer with YOLO") as demo:
+        gr.Markdown("# 🔍 Image Space Analyzer with YOLO")
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                image_input = gr.Image(type="pil", label="📸 Upload Image", height=500)
+                status_output = gr.Textbox(label="📋 Status", interactive=False, max_lines=2)
+            with gr.Column(scale=1):
+                model_status = gr.Textbox(label="🤖 Model Status", value="Loading YOLO model...", interactive=False)
+                
+                gr.Markdown("### 🎯 YOLO Object Detection")
+                confidence_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.1, label="Confidence Threshold")
+                merge_checkbox = gr.Checkbox(label="Merge with existing selections", value=True)
+                detect_btn = gr.Button("🔍 Detect Objects", variant="primary")
+                
+                gr.Markdown("### ✏️ Manual Selection")
+                manual_input = gr.Textbox(label="Selection (x,y,width,height)", placeholder="100,100,200,150", info="Enter coordinates separated by commas")
+                add_manual_btn = gr.Button("➕ Add Manual Selection")
+                
+                gr.Markdown("### 🗂️ Selection Management")
+                clear_all_btn = gr.Button("🗑️ Clear All", variant="secondary")
+                clear_yolo_btn = gr.Button("🗑️ Clear YOLO")
+                clear_manual_btn = gr.Button("🗑️ Clear Manual")
+        
+        analysis_output = gr.Markdown(label="📊 Analysis Results")
+        
+        image_input.upload(process_image_upload, inputs=[image_input], outputs=[image_input, status_output, analysis_output])
+        detect_btn.click(detect_objects, inputs=[image_input, confidence_slider, merge_checkbox], outputs=[image_input, status_output, analysis_output])
+        add_manual_btn.click(add_manual_selection, inputs=[image_input, manual_input], outputs=[image_input, status_output, analysis_output])
+        clear_all_btn.click(clear_all_selections, inputs=[image_input], outputs=[image_input, status_output, analysis_output])
+        clear_yolo_btn.click(clear_yolo_selections, inputs=[image_input], outputs=[image_input, status_output, analysis_output])
+        clear_manual_btn.click(clear_manual_selections, inputs=[image_input], outputs=[image_input, status_output, analysis_output])
+        
+        demo.load(load_yolo_model, outputs=[model_status])
+    
+    return demo
+
+if __name__ == "__main__":
+    print("🚀 Starting Image Space Analyzer on Hugging Face Spaces...")
+    demo = create_interface()
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)