app.py

#!/usr/bin/env python3
"""
GeoVLM with 3D Globe Visualization
Interactive 3D globe that flies to predicted locations
"""

import gradio as gr
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageTextToText
import torch
import re
import json
from dataclasses import dataclass

# ============================================================================
# Geolocation Parser
# ============================================================================

@dataclass(frozen=True)
class Coords:
    lat: float
    lon: float

@dataclass(frozen=True)
class ParsedResponse:
    city: str | None
    region: str | None
    country: str | None
    coords: Coords | None
    raw_text: str
    format_valid: bool

PROMPT_TEMPLATE = (
    "Look at the image and guess the location.\n"
    "Respond with EXACTLY these 5 lines, no extra text:\n"
    "City: <city name>\n"
    "Region: <state or region>\n"
    "Country: <country name or ISO-2 code>\n"
    "Latitude: <number between -90 and 90>\n"
    "Longitude: <number between -180 and 180>\n"
)

KEY_ALIASES = {
    "city": "city", "country": "country", "region": "region",
    "state": "region", "province": "region",
    "latitude": "lat", "lat": "lat",
    "longitude": "lon", "lon": "lon",
}

def parse_response(text: str) -> ParsedResponse:
    """Parse structured 5-line format"""
    parsed = {}
    if not text:
        return ParsedResponse(None, None, None, None, text, False)
    
    key_pattern = re.compile(
        r'^\s*(?:[-*+\u2022]\s*)?(?P<key>[A-Za-z][A-Za-z0-9\s\-/_.]*?)\s*:\s*(?P<value>.+)$'
    )
    
    for line in text.splitlines():
        match = key_pattern.match(line)
        if not match:
            continue
        
        key_raw = match.group("key").strip().lower()
        key_raw = re.sub(r"\s+", " ", key_raw.strip("*_`\"' "))
        canonical = KEY_ALIASES.get(key_raw)
        
        if canonical is None:
            continue
        
        value_raw = match.group("value").strip().strip("`\"' \t")
        value_raw = re.sub(r"^[*_`]+|[*_`]+$", "", value_raw).strip()
        
        if canonical in {"city", "region", "country"}:
            if value_raw and canonical not in parsed:
                parsed[canonical] = value_raw
        elif canonical in {"lat", "lon"}:
            if canonical not in parsed:
                match_num = re.search(r"-?\d+(?:[.,]\d+)?", value_raw)
                if match_num:
                    try:
                        parsed[canonical] = float(match_num.group(0).replace(",", "."))
                    except ValueError:
                        pass
    
    coords = None
    if "lat" in parsed and "lon" in parsed:
        try:
            lat, lon = parsed["lat"], parsed["lon"]
            if -90 <= lat <= 90 and -180 <= lon <= 180:
                coords = Coords(lat=lat, lon=lon)
        except (ValueError, TypeError):
            pass
    
    return ParsedResponse(
        city=parsed.get("city"),
        region=parsed.get("region"),
        country=parsed.get("country"),
        coords=coords,
        raw_text=text,
        format_valid=bool(len(parsed) >= 2),
    )

# ============================================================================
# Model Setup
# ============================================================================

model = None
processor = None
MODEL_NAME = "Qwen/Qwen2-VL-2B-Instruct"

def load_model():
    global model, processor
    if model is None:
        print(f"Loading model: {MODEL_NAME}")
        processor = AutoProcessor.from_pretrained(MODEL_NAME)
        model = AutoModelForImageTextToText.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else "cpu"
        )
        print("Model loaded!")

def predict_location(image):
    """Predict geolocation and return globe visualization data"""
    if image is None:
        return "Please upload an image.", "", ""
    
    load_model()
    
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image).convert("RGB")
    else:
        image = image.convert("RGB")
    
    messages = [{
        "role": "user",
        "content": [
            {"type": "image"},
            {"type": "text", "text": PROMPT_TEMPLATE}
        ]
    }]
    
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text], images=[image], return_tensors="pt", padding=True)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=256, do_sample=False)
    
    generated_ids = output_ids[0][inputs['input_ids'].shape[1]:]
    response = processor.decode(generated_ids, skip_special_tokens=True).strip()
    parsed = parse_response(response)
    
    # Format output
    output = f"""
## 🤖 AI Prediction

**📍 Location Details:**
- **City:** {parsed.city or "Unknown"}
- **Region:** {parsed.region or "Unknown"}
- **Country:** {parsed.country or "Unknown"}
- **Coordinates:** {f"{parsed.coords.lat:.6f}°, {parsed.coords.lon:.6f}°" if parsed.coords else "Not found"}

---

## 🔍 Raw Response:
```
{response}
```
"""
    
    # Create globe HTML
    globe_html = create_globe_html(parsed) if parsed.coords else "<div style='text-align:center; padding:50px; color:#666;'>No coordinates found</div>"
    
    # Create info card
    info_html = create_info_card(parsed)
    
    return output, globe_html, info_html

def create_globe_html(parsed: ParsedResponse) -> str:
    """Create Three.js globe visualization with day/night toggle and country borders"""
    if not parsed.coords:
        return ""
    
    lat, lon = parsed.coords.lat, parsed.coords.lon
    
    html = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <style>
            body {{ margin: 0; padding: 0; overflow: hidden; background: #000; position: relative; }}
            #globeViz {{ width: 100%; height: 600px; }}
            .location-label {{
                color: white;
                font-size: 16px;
                font-family: Arial, sans-serif;
                background: rgba(0,0,0,0.7);
                padding: 8px 12px;
                border-radius: 4px;
                pointer-events: none;
            }}
            .controls {{
                position: absolute;
                top: 20px;
                right: 20px;
                z-index: 100;
                display: flex;
                gap: 10px;
            }}
            .control-btn {{
                background: rgba(255,255,255,0.9);
                border: none;
                padding: 10px 16px;
                border-radius: 6px;
                cursor: pointer;
                font-weight: bold;
                font-size: 14px;
                transition: all 0.3s;
                box-shadow: 0 2px 8px rgba(0,0,0,0.3);
            }}
            .control-btn:hover {{
                background: white;
                transform: translateY(-2px);
                box-shadow: 0 4px 12px rgba(0,0,0,0.4);
            }}
            .control-btn.active {{
                background: #667eea;
                color: white;
            }}
        </style>
    </head>
    <body>
        <div class="controls">
            <button class="control-btn active" id="dayBtn" onclick="setDayMode()">☀️ Day</button>
            <button class="control-btn" id="nightBtn" onclick="setNightMode()">🌙 Night</button>
            <button class="control-btn" id="bordersBtn" onclick="toggleBorders()">🗺️ Borders</button>
        </div>
        <div id="globeViz"></div>
        
        <script src="//unpkg.com/globe.gl"></script>
        <script>
            let showBorders = false;
            let currentMode = 'day';
            
            const myGlobe = Globe()
                .globeImageUrl('//unpkg.com/three-globe/example/img/earth-blue-marble.jpg')
                .bumpImageUrl('//unpkg.com/three-globe/example/img/earth-topology.png')
                .backgroundImageUrl('//unpkg.com/three-globe/example/img/night-sky.png')
                .pointOfView({{ lat: {lat}, lng: {lon}, altitude: 2.5 }}, 0)
                .atmosphereColor('lightskyblue')
                .atmosphereAltitude(0.15)
                (document.getElementById('globeViz'));

            // Load country borders
            fetch('//unpkg.com/world-atlas/countries-50m.json')
                .then(res => res.json())
                .then(countries => {{
                    window.countriesData = countries;
                }});

            // Add marker point
            const markerData = [{{
                lat: {lat},
                lng: {lon},
                size: 0.5,
                color: '#ff4444',
                label: '{parsed.city or "Location"}',
                city: '{parsed.city or "Unknown"}',
                region: '{parsed.region or "Unknown"}',
                country: '{parsed.country or "Unknown"}'
            }}];

            myGlobe
                .pointsData(markerData)
                .pointAltitude('size')
                .pointColor('color')
                .pointRadius(0.6)
                .pointLabel(d => `
                    <div class="location-label">
                        <b>${{d.city}}</b><br/>
                        ${{d.region}}, ${{d.country}}<br/>
                        ${{d.lat.toFixed(4)}}°, ${{d.lng.toFixed(4)}}°
                    </div>
                `);

            // Animate to location
            myGlobe.pointOfView({{ lat: {lat}, lng: {lon}, altitude: 1.5 }}, 3000);

            // Auto-rotate
            myGlobe.controls().autoRotate = true;
            myGlobe.controls().autoRotateSpeed = 0.3;

            // Add pulsing ring animation
            const ringData = [{{
                lat: {lat},
                lng: {lon},
                maxR: 10,
                propagationSpeed: 2,
                repeatPeriod: 1500
            }}];

            myGlobe
                .ringsData(ringData)
                .ringColor(() => 'rgba(255,68,68,0.5)')
                .ringMaxRadius('maxR')
                .ringPropagationSpeed('propagationSpeed')
                .ringRepeatPeriod('repeatPeriod');

            // Add arcs for visual effect
            const arcData = [{{
                startLat: {lat},
                startLng: {lon},
                endLat: {lat + 10},
                endLng: {lon + 10},
                color: ['rgba(255,68,68,0.4)', 'rgba(255,68,68,0.1)']
            }}];

            myGlobe
                .arcsData(arcData)
                .arcColor('color')
                .arcDashLength(0.4)
                .arcDashGap(0.2)
                .arcDashAnimateTime(2000)
                .arcStroke(0.5);

            // Mode switching functions
            function setDayMode() {{
                currentMode = 'day';
                myGlobe
                    .globeImageUrl('//unpkg.com/three-globe/example/img/earth-blue-marble.jpg')
                    .bumpImageUrl('//unpkg.com/three-globe/example/img/earth-topology.png');
                
                document.getElementById('dayBtn').classList.add('active');
                document.getElementById('nightBtn').classList.remove('active');
            }}

            function setNightMode() {{
                currentMode = 'night';
                myGlobe
                    .globeImageUrl('//unpkg.com/three-globe/example/img/earth-night.jpg')
                    .bumpImageUrl('//unpkg.com/three-globe/example/img/earth-topology.png');
                
                document.getElementById('nightBtn').classList.add('active');
                document.getElementById('dayBtn').classList.remove('active');
            }}

            function toggleBorders() {{
                showBorders = !showBorders;
                const btn = document.getElementById('bordersBtn');
                
                if (showBorders && window.countriesData) {{
                    const countries = topojson.feature(window.countriesData, window.countriesData.objects.countries);
                    myGlobe
                        .polygonsData(countries.features)
                        .polygonAltitude(0.01)
                        .polygonCapColor(() => 'rgba(200, 200, 200, 0.1)')
                        .polygonSideColor(() => 'rgba(200, 200, 200, 0.05)')
                        .polygonStrokeColor(() => '#ffffff')
                        .polygonLabel(({{ properties: d }}) => `
                            <div class="location-label">
                                <b>${{d.name}}</b>
                            </div>
                        `);
                    btn.classList.add('active');
                }} else {{
                    myGlobe.polygonsData([]);
                    btn.classList.remove('active');
                }}
            }}
        </script>
        <script src="//unpkg.com/topojson-client"></script>
    </body>
    </html>
    """
    return html

def create_info_card(parsed: ParsedResponse) -> str:
    """Create information card with details"""
    if not parsed.coords:
        return ""
    
    lat, lon = parsed.coords.lat, parsed.coords.lon
    
    html = f"""
    <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                border-radius: 12px; padding: 24px; color: white; margin-top: 20px;">
        <h2 style="margin: 0 0 16px 0; font-size: 24px;">📍 Predicted Location</h2>
        
        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 16px; margin-bottom: 20px;">
            <div style="background: rgba(255,255,255,0.1); padding: 12px; border-radius: 8px;">
                <div style="font-size: 12px; opacity: 0.8;">City</div>
                <div style="font-size: 18px; font-weight: bold;">{parsed.city or "Unknown"}</div>
            </div>
            <div style="background: rgba(255,255,255,0.1); padding: 12px; border-radius: 8px;">
                <div style="font-size: 12px; opacity: 0.8;">Region</div>
                <div style="font-size: 18px; font-weight: bold;">{parsed.region or "Unknown"}</div>
            </div>
            <div style="background: rgba(255,255,255,0.1); padding: 12px; border-radius: 8px;">
                <div style="font-size: 12px; opacity: 0.8;">Country</div>
                <div style="font-size: 18px; font-weight: bold;">{parsed.country or "Unknown"}</div>
            </div>
            <div style="background: rgba(255,255,255,0.1); padding: 12px; border-radius: 8px;">
                <div style="font-size: 12px; opacity: 0.8;">Coordinates</div>
                <div style="font-size: 14px; font-weight: bold;">{lat:.4f}°, {lon:.4f}°</div>
            </div>
        </div>
        
        <div style="display: flex; gap: 12px; flex-wrap: wrap;">
            <a href="https://www.google.com/maps?q={lat},{lon}" target="_blank" 
               style="background: #4285f4; color: white; padding: 10px 20px; 
                      border-radius: 6px; text-decoration: none; font-weight: bold;">
                🗺️ Google Maps
            </a>
            <a href="https://www.openstreetmap.org/?mlat={lat}&mlon={lon}#map=12/{lat}/{lon}" target="_blank"
               style="background: #7ebc6f; color: white; padding: 10px 20px; 
                      border-radius: 6px; text-decoration: none; font-weight: bold;">
                🌍 OpenStreetMap
            </a>
            <a href="https://www.google.com/search?q={parsed.city}+{parsed.country}" target="_blank"
               style="background: #ea4335; color: white; padding: 10px 20px; 
                      border-radius: 6px; text-decoration: none; font-weight: bold;">
                🔍 Learn More
            </a>
        </div>
    </div>
    """
    return html

# ============================================================================
# Gradio Interface
# ============================================================================

with gr.Blocks(title="GeoVLM - 3D Globe", theme=gr.themes.Soft(), css="""
    .gradio-container {max-width: 1400px !important;}
    .globe-container {height: 600px !important;}
""") as demo:
    
    gr.Markdown("""
    # 🌍 GeoVLM - AI Geolocation with 3D Globe
    
    Upload any image and watch the AI predict its location on an interactive 3D globe!
    
    **Powered by:** [vlm-gym](https://github.com/sdan/vlm-gym) | Vision-Language Models | Three.js Globe
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil",
                label="📸 Upload Image",
                height=400
            )
            
            predict_btn = gr.Button(
                "🔍 Analyze & Locate", 
                variant="primary", 
                size="lg"
            )
            
            gr.Markdown("""
            ### 💡 Tips:
            - Outdoor images work best
            - Street views are ideal  
            - Landmarks help accuracy
            - Clear, well-lit photos
            
            ### 🎯 Features:
            - 3D interactive globe
            - Flies to predicted location
            - Pulsing marker animation
            - Auto-rotating globe
            """)
        
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("🌐 3D Globe"):
                    globe_output = gr.HTML(
                        label="Interactive Globe",
                        elem_classes=["globe-container"]
                    )
                
                with gr.Tab("📊 Details"):
                    info_output = gr.HTML(label="Location Info")
                    output_text = gr.Markdown(label="Analysis")
    
    gr.Markdown("""
    ---
    
    ### 🎮 How It Works:
    
    1. **Upload** any image with visible location clues
    2. **AI analyzes** architecture, vegetation, signs, landscape
    3. **Globe flies** to the predicted location in 3D
    4. **Explore** the area with interactive controls
    
    ### 🔬 Technology:
    - **Vision Model:** Qwen2-VL-2B-Instruct
    - **Training:** Reinforcement learning on 5M geotagged images
    - **Visualization:** Three.js Globe.GL
    - **Dataset:** OSV5M (OpenStreetView 5M)
    
    ### 🚀 Use Cases:
    - **OSINT Research** - Verify photo locations
    - **Education** - Learn world geography
    - **Travel** - Discover new places
    - **Training** - Practice geolocation skills
    
    ---
    
    Built with ❤️ by AceXRoux | [GitHub](https://github.com/sdan/vlm-gym) | [LinkedIn](https://linkedin.com/in/your-profile)
    """)
    
    predict_btn.click(
        fn=predict_location,
        inputs=image_input,
        outputs=[output_text, globe_output, info_output]
    )

if __name__ == "__main__":
    print("🌍 Starting GeoVLM with 3D Globe...")
    load_model()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )