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