app.py

# app.py

import gradio as gr
import numpy as np
import cv2
from PIL import Image, ImageDraw, ImageFont
import json
import time
from ultralytics import YOLO
import torch
import os

# Fix Ultralytics config path for Hugging Face Spaces (read-only /home/user)
os.environ['YOLO_CONFIG_DIR'] = '/tmp'

# Global variables to maintain state
model = None

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,
            'x': x,
            'y': y,
            'width': width,
            'height': height,
            'type': selection_type,
            'label': label,
            'confidence': confidence,
            'original_aspect_ratio': width / (height + 1e-6)
        }
        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
        return total_selected_area, total_image_area

# Initialize selection manager
selection_manager = SelectionManager()

def load_yolo_model():
    global model
    try:
        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)}"

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  # ms
        
        if not merge_with_existing:
            selection_manager.clear_type('yolo')
        
        detections_added = 0
        for result in results:
            boxes = result.boxes
            if boxes is not None:
                for box in boxes:
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    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),
                        width=int(x2 - x1),
                        height=int(y2 - y1),
                        selection_type='yolo',
                        label=f"{class_name} ({confidence:.2f})",
                        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)

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']
        
        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

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

**Total Image Area:** {total_image_area:,} px²  
**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)}  

### 📋 Selection Details:
"""
    for i, selection in enumerate(selection_manager.selections, 1):
        area = selection['width'] * selection['height']
        analysis += f"""
**{i}.** {selection['label'] or f"Selection {selection['id']}"}  
- Type: {selection['type'].upper()}  
- Area: {area:,} px²  
- Dimensions: {selection['width']}×{selection['height']}  
"""
        if selection['confidence'] is not None:
            analysis += f"- Confidence: {selection['confidence']:.2f}\n"
    
    return analysis

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 {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)

def clear_all_selections(image):
    selection_manager.clear_all()
    if image is not None:
        return image.copy(), "✅ All selections cleared", get_analysis_text(image)
    return None, "✅ All selections cleared", ""

def clear_yolo_selections(image):
    selection_manager.clear_type('yolo')
    if image is not None:
        annotated_image = draw_selections(image)
        return annotated_image, "✅ YOLO detections cleared", get_analysis_text(image)
    return None, "✅ YOLO detections cleared", ""

def clear_manual_selections(image):
    selection_manager.clear_type('manual')
    if image is not None:
        annotated_image = draw_selections(image)
        return annotated_image, "✅ Manual selections cleared", get_analysis_text(image)
    return None, "✅ Manual selections cleared", ""

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
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)