Upload demo_oculus_unified.py with huggingface_hub

demo_oculus_unified.py

@@ -0,0 +1,263 @@
+#!/usr/bin/env python3
+"""
+Oculus 0.2 Unified Demo
+
+Demonstrates all features of the unified Oculus model:
+- Text mode (captioning, VQA)
+- Point mode (counting objects)
+- Box mode (detection with bounding boxes)
+- Polygon mode (segmentation)
+- Optional reasoning with thinking traces
+- Focus system for fine-grained perception
+"""
+
+import os
+import sys
+import requests
+from pathlib import Path
+from io import BytesIO
+
+from PIL import Image
+import torch
+
+# Add parent to path
+sys.path.insert(0, str(Path(__file__).parent))
+
+from oculus_unified_model import OculusForConditionalGeneration, OculusConfig
+
+
+def download_image(url: str) -> Image.Image:
+    """Download image from URL."""
+    headers = {'User-Agent': 'Mozilla/5.0'}
+    response = requests.get(url, headers=headers, timeout=10)
+    response.raise_for_status()
+    return Image.open(BytesIO(response.content)).convert('RGB')
+
+
+def print_header(title: str):
+    print("\n" + "=" * 70)
+    print(f"🔮 {title}")
+    print("=" * 70)
+
+
+def print_section(title: str):
+    print(f"\n{'─' * 70}")
+    print(f"   {title}")
+    print(f"{'─' * 70}")
+
+
+def demo():
+    print_header("OCULUS 0.2 UNIFIED MODEL DEMO")
+    
+    # ================================================================
+    # Load Model
+    # ================================================================
+    print("\n[1] Loading Oculus Model...")
+    
+    # Check if we have trained weights
+    weights_path = Path(__file__).parent / "checkpoints" / "oculus_coco" / "final"
+    
+    if weights_path.exists():
+        print(f"    Found trained weights at: {weights_path}")
+        model = OculusForConditionalGeneration.from_pretrained(weights_path)
+    else:
+        print("    Using default configuration")
+        config = OculusConfig(
+            reasoning_enabled=True,
+            enable_focus=True,
+        )
+        model = OculusForConditionalGeneration(config)
+    
+    print("    ✓ Model loaded!")
+    
+    # ================================================================
+    # Test Images
+    # ================================================================
+    test_images = [
+        {
+            "name": "Cat on Couch",
+            "url": "https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg"
+        },
+        {
+            "name": "Golden Gate Bridge",
+            "url": "https://upload.wikimedia.org/wikipedia/commons/thumb/0/0c/GoldenGateBridge-001.jpg/1200px-GoldenGateBridge-001.jpg"
+        },
+    ]
+    
+    for test in test_images:
+        print_header(f"Testing: {test['name']}")
+        
+        try:
+            print("\n[Downloading image...]")
+            image = download_image(test["url"])
+            print(f"    Image size: {image.size}")
+            
+            # ========================================================
+            # Mode 1: TEXT (Captioning)
+            # ========================================================
+            print_section("📝 TEXT MODE - Captioning")
+            
+            output = model.generate(
+                image=image,
+                prompt="Describe this image in detail",
+                mode="text",
+                think=False
+            )
+            
+            print(f"    Caption: \"{output.text}\"")
+            
+            # ========================================================
+            # Mode 2: TEXT with Reasoning
+            # ========================================================
+            print_section("🧠 TEXT MODE - With Reasoning")
+            
+            output = model.generate(
+                image=image,
+                prompt="What is the main subject of this image?",
+                mode="text",
+                think=True  # Enable thinking traces
+            )
+            
+            if output.thinking_trace:
+                print(f"    💭 Thinking: {output.thinking_trace[:200]}...")
+            print(f"    Answer: \"{output.text}\"")
+            
+            # ========================================================
+            # Mode 3: TEXT (VQA)
+            # ========================================================
+            print_section("❓ TEXT MODE - VQA")
+            
+            questions = [
+                "What colors are visible in this image?",
+                "Is this indoors or outdoors?",
+            ]
+            
+            for q in questions:
+                output = model.generate(
+                    image=image,
+                    prompt=q,
+                    mode="text"
+                )
+                print(f"    Q: {q}")
+                print(f"    A: {output.text}")
+            
+            # ========================================================
+            # Mode 4: POINT (Counting)
+            # ========================================================
+            print_section("📍 POINT MODE - Object Counting")
+            
+            output = model.generate(
+                image=image,
+                prompt="Find objects",
+                mode="point"
+            )
+            
+            print(f"    Detected {len(output.points)} points")
+            for i, (pt, label, conf) in enumerate(zip(
+                output.points[:5],
+                output.labels[:5],
+                output.confidences[:5]
+            )):
+                print(f"      Point {i+1}: {pt} (class={label}, conf={conf:.2f})")
+            
+            # ========================================================
+            # Mode 5: BOX (Detection)
+            # ========================================================
+            print_section("📦 BOX MODE - Object Detection")
+            
+            output = model.generate(
+                image=image,
+                prompt="Detect all objects",
+                mode="box"
+            )
+            
+            print(f"    Detected {len(output.boxes)} boxes")
+            for i, (box, label, conf) in enumerate(zip(
+                output.boxes[:5],
+                output.labels[:5],
+                output.confidences[:5]
+            )):
+                print(f"      Box {i+1}: {[f'{b:.2f}' for b in box]} (class={label}, conf={conf:.2f})")
+            
+            # ========================================================
+            # Mode 6: POLYGON (Segmentation)
+            # ========================================================
+            print_section("🔷 POLYGON MODE - Segmentation")
+            
+            output = model.generate(
+                image=image,
+                prompt="Segment the scene",
+                mode="polygon"
+            )
+            
+            print(f"    Segmentation mask shape: {output.mask.shape if output.mask is not None else 'N/A'}")
+            print(f"    Detected {len(output.polygons)} regions")
+            for i, (poly, label) in enumerate(zip(
+                output.polygons[:3],
+                output.labels[:3]
+            )):
+                print(f"      Region {i+1}: class={label}, vertices={len(poly)}")
+            
+            print("\n    ✅ All modes successful!")
+            
+        except Exception as e:
+            print(f"\n    ❌ Error: {e}")
+            import traceback
+            traceback.print_exc()
+    
+    # ================================================================
+    # Summary
+    # ================================================================
+    print_header("DEMO COMPLETE")
+    
+    print("""
+    Oculus 0.2 supports:
+    
+    📝 TEXT MODE
+       - Image captioning
+       - Visual question answering
+       - With optional reasoning traces
+    
+    📍 POINT MODE
+       - Object counting
+       - Point localization
+    
+    📦 BOX MODE
+       - Object detection
+       - Bounding box prediction
+    
+    🔷 POLYGON MODE
+       - Semantic segmentation
+       - Instance segmentation
+    
+    🧠 REASONING
+       - Optional thinking traces
+       - Multi-step reasoning
+    
+    🔍 FOCUS SYSTEM
+       - Zoom & crop for fine-grained perception
+       - Automatic region detection
+    
+    Usage:
+    ```python
+    from oculus_unified_model import OculusForConditionalGeneration
+    
+    model = OculusForConditionalGeneration.from_pretrained("./checkpoints/oculus_coco/final")
+    
+    # Caption
+    output = model.generate(image, mode="text", prompt="Describe this")
+    
+    # VQA with reasoning
+    output = model.generate(image, mode="text", prompt="What color is it?", think=True)
+    
+    # Detection
+    output = model.generate(image, mode="box", prompt="Find cars")
+    
+    # Segmentation
+    output = model.generate(image, mode="polygon")
+    ```
+    """)
+
+
+if __name__ == "__main__":
+    demo()