app.py

import gradio as gr
from transformers import pipeline, ViTImageProcessor, ViTForImageClassification, YolosImageProcessor, YolosForObjectDetection
from PIL import Image, ImageDraw
import torch
import numpy as np

class UniversalImageClassifier:
    def __init__(self):
        print("🔄 Cargando clasificador de imágenes ViT...")
        self.model_name = "google/vit-base-patch16-224"
        self.processor = ViTImageProcessor.from_pretrained(self.model_name)
        self.model = ViTForImageClassification.from_pretrained(self.model_name)
        
        self.classifier = pipeline(
            "image-classification",
            model=self.model,
            feature_extractor=self.processor,
            device=-1  # CPU
        )
        
        print("✅ Clasificador ViT cargado!")
        
        self.category_mappings = {
            'egyptian_cat': '🐱 Gato Egipcio',
            'tabby': '🐱 Gato Atigrado', 
            'tiger_cat': '🐱 Gato Tiger',
            'golden_retriever': '🐕 Golden Retriever',
            'german_shepherd': '🐕 Pastor Alemán',
            'beagle': '🐕 Beagle',
            'sports_car': '🏎️ Auto Deportivo',
            'convertible': '🚗 Convertible',
            'motorcycle': '🏍️ Motocicleta',
            'bicycle': '🚲 Bicicleta',
            'airplane': '✈️ Avión',
            'pizza': '🍕 Pizza',
            'hamburger': '🍔 Hamburguesa',
            'hot_dog': '🌭 Hot Dog',
            'ice_cream': '🍦 Helado',
            'laptop': '💻 Laptop',
            'cellular_telephone': '📱 Teléfono Móvil',
            'television': '📺 Televisión',
            'daisy': '🌼 Margarita',
            'rose': '🌹 Rosa',
            'sunflower': '🌻 Girasol',
        }
    
    def classify_image(self, image):
        try:
            results = self.classifier(image)
            predictions = []
            for result in results[:5]:
                label = result['label']
                confidence = result['score'] * 100
                display_label = self.category_mappings.get(label, f"🔍 {label.replace('_', ' ').title()}")
                
                predictions.append({
                    'label': display_label,
                    'original_label': label,
                    'confidence': confidence
                })
            
            return predictions
        except Exception as e:
            return [{'label': f'Error: {str(e)}', 'confidence': 0}]

class ObjectDetector:
    def __init__(self):
        print("🔄 Cargando detector de objetos YOLOS...")
        self.model_name = "hustvl/yolos-tiny"
        self.processor = YolosImageProcessor.from_pretrained(self.model_name)
        self.model = YolosForObjectDetection.from_pretrained(self.model_name)
        
        print("✅ Detector YOLOS cargado!")
        
        self.class_mappings = {
            'person': '👤 Persona',
            'bicycle': '🚲 Bicicleta', 
            'car': '🚗 Auto',
            'motorcycle': '🏍️ Motocicleta',
            'airplane': '✈️ Avión',
            'bus': '🚌 Autobús',
            'train': '🚂 Tren',
            'truck': '🚛 Camión',
            'boat': '⛵ Barco',
            'traffic light': '🚦 Semáforo',
            'bird': '🐦 Pájaro',
            'cat': '🐱 Gato',
            'dog': '🐕 Perro',
            'horse': '🐎 Caballo',
            'elephant': '🐘 Elefante',
            'backpack': '🎒 Mochila',
            'umbrella': '☂️ Paraguas',
            'bottle': '🍼 Botella',
            'cup': '☕ Taza',
            'banana': '🍌 Plátano',
            'apple': '🍎 Manzana',
            'pizza': '🍕 Pizza',
            'chair': '🪑 Silla',
            'couch': '🛋️ Sofá',
            'bed': '🛏️ Cama',
            'tv': '📺 Televisión',
            'laptop': '💻 Laptop',
            'mouse': '🖱️ Mouse',
            'cell phone': '📱 Teléfono',
            'book': '📚 Libro',
            'clock': '🕐 Reloj'
        }
    
    def detect_objects(self, image):
        try:
            inputs = self.processor(images=image, return_tensors="pt")
            outputs = self.model(**inputs)
            
            target_sizes = torch.tensor([image.size[::-1]])
            results = self.processor.post_process_object_detection(
                outputs, target_sizes=target_sizes, threshold=0.3
            )[0]
            
            detections = []
            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
                if score > 0.3:
                    label_name = self.model.config.id2label[label.item()]
                    display_name = self.class_mappings.get(label_name, f"🔍 {label_name}")
                    
                    detections.append({
                        'label': display_name,
                        'original_label': label_name,
                        'confidence': score.item() * 100,
                        'box': box.tolist()
                    })
            
            return detections
        except Exception as e:
            return [{'label': f'Error: {str(e)}', 'confidence': 0, 'box': [0, 0, 0, 0]}]
    
    def draw_detections(self, image, detections):
        try:
            annotated_image = image.copy()
            draw = ImageDraw.Draw(annotated_image)
            
            colors = ['red', 'blue', 'green', 'yellow', 'purple', 'orange', 'cyan', 'magenta']
            
            for i, detection in enumerate(detections):
                if detection['confidence'] > 30:
                    box = detection['box']
                    label = detection['label']
                    confidence = detection['confidence']
                    
                    xmin, ymin, xmax, ymax = box
                    color = colors[i % len(colors)]
                    
                    draw.rectangle([(xmin, ymin), (xmax, ymax)], outline=color, width=3)
                    
                    text = f"{label} ({confidence:.1f}%)"
                    try:
                        text_bbox = draw.textbbox((0, 0), text)
                        text_width = text_bbox[2] - text_bbox[0]
                        text_height = text_bbox[3] - text_bbox[1]
                        
                        draw.rectangle(
                            [(xmin, ymin - text_height - 4), (xmin + text_width + 4, ymin)],
                            fill=color
                        )
                        draw.text((xmin + 2, ymin - text_height - 2), text, fill='white')
                    except:
                        draw.text((xmin, ymin - 20), text, fill=color)
            
            return annotated_image
        except Exception as e:
            return image

# Inicializar modelos
print("🚀 Inicializando modelos de IA...")
classifier = UniversalImageClassifier()
detector = ObjectDetector()
print("✅ ¡Todos los modelos listos!")

def classify_image_complete(image):
    if image is None:
        return "❌ Por favor sube una imagen para clasificar"
    
    try:
        predictions = classifier.classify_image(image)
        
        if not predictions or predictions[0]['confidence'] == 0:
            return "❌ No se pudo clasificar la imagen"
        
        dominant = predictions[0]
        
        report = f"""# 🔍 Clasificación de Imagen
## 🎯 Predicción Principal:
### {dominant['label']} 
**Confianza:** {dominant['confidence']:.1f}%
## 📊 Top 5 Predicciones:
"""
        
        for i, pred in enumerate(predictions, 1):
            bar = "█" * int(pred['confidence'] / 10) + "░" * (10 - int(pred['confidence'] / 10))
            report += f"\n**{i}.** {pred['label']}\n{bar} {pred['confidence']:.1f}%"
        
        # Análisis de confianza
        confidence = dominant['confidence']
        if confidence > 80:
            level = "🟢 Muy Alta"
        elif confidence > 60:
            level = "🟡 Alta"
        elif confidence > 40:
            level = "🟠 Moderada"
        else:
            level = "🔴 Baja"
        
        report += f"\n\n## 🎚️ Confianza: {level}"
        report += f"\n\n*Clasificación con Vision Transformer (ViT)*"
        
        return report
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

def detect_objects_complete(image):
    if image is None:
        return "❌ Por favor sube una imagen para detectar objetos", None
    
    try:
        detections = detector.detect_objects(image)
        
        if not detections or detections[0]['confidence'] == 0:
            return "❌ No se detectaron objetos", image
        
        annotated_image = detector.draw_detections(image, detections)
        
        object_counts = {}
        for detection in detections:
            if detection['confidence'] > 30:
                label = detection['original_label']
                object_counts[label] = object_counts.get(label, 0) + 1
        
        total_objects = len(detections)
        unique_objects = len(object_counts)
        
        report = f"""# 🎯 Detección de Objetos
## 📊 Resumen:
- **Objetos detectados:** {total_objects}
- **Tipos únicos:** {unique_objects}
- **Confianza promedio:** {np.mean([d['confidence'] for d in detections]):.1f}%
## 🔍 Objetos Encontrados:
"""
        
        sorted_detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)
        
        for i, detection in enumerate(sorted_detections[:10], 1):  # Top 10
            label = detection['label']
            confidence = detection['confidence']
            bar = "█" * int(confidence / 10) + "░" * (10 - int(confidence / 10))
            
            report += f"\n**{i}.** {label}\n{bar} {confidence:.1f}%"
        
        report += f"\n\n## 📈 Conteo por Tipo:"
        for obj_type, count in sorted(object_counts.items(), key=lambda x: x[1], reverse=True)[:5]:
            display_name = detector.class_mappings.get(obj_type, obj_type)
            report += f"\n- {display_name}: **{count}**"
        
        report += f"\n\n*Detección con YOLOS (You Only Look Once)*"
        
        return report, annotated_image
        
    except Exception as e:
        return f"❌ Error: {str(e)}", None

# Interfaz Gradio
with gr.Blocks(title="🤖 Analizador Visual Universal con IA") as demo:
    
    gr.Markdown("""
    # 🤖 Analizador Visual Universal con IA
    
    **Dos poderosos modelos de IA en una sola aplicación**
    
    🔍 **Clasificador Universal:** Identifica QUÉ ES (1000+ categorías)  
    🎯 **Detector de Objetos:** Encuentra DÓNDE ESTÁN (80+ objetos)
    
    ✨ **Modelos incluidos:**
    - 🧠 Vision Transformer (ViT) de Google
    - 🎯 YOLOS (You Only Look Once)
    - 💻 100% optimizado para CPU
    """)
    
    with gr.Tabs():
        # Tab 1: Clasificador
        with gr.Tab("🔍 Clasificador Universal"):
            gr.Markdown("""
            ### Identifica automáticamente el contenido principal de tu imagen
            **Perfecto para:** Catalogar fotos, identificar objetos desconocidos, análisis de contenido
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    classify_input = gr.Image(
                        label="📸 Sube tu imagen",
                        type="pil"
                    )
                    classify_btn = gr.Button(
                        "🚀 Clasificar Imagen", 
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=2):
                    classify_output = gr.Markdown(
                        label="📋 Resultado de Clasificación"
                    )
            
            # gr.Examples(
            #     examples=[
            #         "https://images.unsplash.com/photo-1574158622682-e40e69881006?w=300",  # Gato
            #         "https://images.unsplash.com/photo-1552053831-71594a27632d?w=300",  # Perro
            #         "https://images.unsplash.com/photo-1565299624946-b28f40a0ca4b?w=300"   # Pizza
            #     ],
            #     inputs=[classify_input],
            #     label="🖼️ Ejemplos para probar"
            # )
        
        # Tab 2: Detector de Objetos
        with gr.Tab("🎯 Detector de Objetos"):
            gr.Markdown("""
            ### Encuentra y localiza múltiples objetos en tu imagen
            **Perfecto para:** Análisis de escenas, inventarios visuales, seguridad
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    detect_input = gr.Image(
                        label="📸 Sube tu imagen",
                        type="pil"
                    )
                    detect_btn = gr.Button(
                        "🎯 Detectar Objetos", 
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=1):
                    detect_output = gr.Markdown(
                        label="📋 Objetos Detectados"
                    )
                    detect_image_output = gr.Image(
                        label="🎯 Imagen Anotada",
                        type="pil"
                    )
    
    # Eventos
    classify_btn.click(
        classify_image_complete,
        inputs=[classify_input],
        outputs=[classify_output]
    )
    
    detect_btn.click(
        detect_objects_complete,
        inputs=[detect_input],
        outputs=[detect_output, detect_image_output]
    )
    
    gr.Markdown("""
    ---
    ### 💡 Consejos para mejores resultados:
    - Usa imágenes claras y bien iluminadas
    - Centra el objeto principal para clasificación
    - Para detección, incluye múltiples objetos en la escena
    - Resolución mínima recomendada: 224x224 píxeles
    
    **🚀 Powered by Hugging Face Transformers**
    """)

if __name__ == "__main__":
    demo.launch()