app.py

import gradio as gr
import spaces
import torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
import os
from datetime import datetime
from PIL import Image
import requests
from io import BytesIO

# Model setup
processor = AutoProcessor.from_pretrained("deepguess/weather-vlm-qwen2.5-7b", trust_remote_code=True)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "deepguess/weather-vlm-qwen2.5-7b",
    torch_dtype=torch.float16,
    trust_remote_code=True
)
model.eval()

# Title and description
TITLE = "🌦️ Weather Analysis VLM (Qwen2.5-VL-7B Fine-tuned)"

DESCRIPTION = """
## Advanced Weather Image Analysis

This model specializes in analyzing weather data including:
- **Model Outputs**: GFS, HRRR, ECMWF, NAM analysis
- **Soundings**: Skew-T diagrams, hodographs, SHARPpy analysis
- **Observations**: Surface obs, satellite, radar imagery
- **Forecasts**: Deterministic and ensemble model outputs
- **Severe Weather**: Convective parameters, SPC outlooks

### ⚠️ Disclaimer
**For educational and research purposes only. Not for operational forecasting.**
"""

# Enhanced prompts based on your data categories
PROMPT_TEMPLATES = {
    "Quick Analysis": "Describe the weather in this image.",
    "Model Output": "Analyze this model output. What patterns and features are shown?",
    "Sounding Analysis": "Analyze this sounding. Discuss stability, shear, and severe potential.",
    "Radar/Satellite": "Describe the features in this radar or satellite image.",
    "Severe Weather": "Assess severe weather potential based on this image.",
    "Technical Deep Dive": "Provide detailed technical analysis including parameters and meteorological significance.",
    "Forecast Discussion": "Based on this image, what weather evolution is expected?",
    "Pattern Recognition": "Identify synoptic patterns, jet streaks, troughs, ridges, and fronts.",
    "Ensemble Analysis": "Analyze ensemble spread, uncertainty, and most likely scenarios.",
    "Winter Weather": "Analyze precipitation type, accumulation potential, and impacts.",
}

# System prompts for different analysis modes
SYSTEM_PROMPTS = {
    "technical": """You are an expert meteorologist providing technical analysis. Focus on:
- Specific parameter values and thresholds
- Physical processes and dynamics
- Pattern recognition and anomalies
- Forecast confidence and uncertainty
Use technical terminology appropriately.""",
    
    "educational": """You are a meteorology instructor. Explain concepts clearly while maintaining accuracy. 
Point out key features and explain their significance. Use some technical terms but define them.""",
    
    "operational": """You are providing a weather briefing. Focus on:
- Current conditions and trends
- Expected evolution
- Impacts and hazards
- Timing of changes
Be concise but thorough.""",
    
    "research": """You are analyzing meteorological data for research purposes. Discuss:
- Interesting features or anomalies
- Comparison to climatology
- Physical mechanisms
- Uncertainty quantification"""
}

# Analysis mode descriptions
MODE_INFO = {
    "technical": "Detailed technical analysis for meteorologists",
    "educational": "Clear explanations for learning",
    "operational": "Focused briefing style",
    "research": "In-depth research perspective"
}

# Example URLs for different weather data types
EXAMPLE_URLS = {
    "SPC Convective Outlook": "https://www.spc.noaa.gov/products/outlook/day1otlk.gif",
    "SPC Mesoanalysis (Surface)": "https://inside.nssl.noaa.gov/ewp/wp-content/uploads/sites/22/2019/05/19ZSPCCape.png",
    "College of DuPage Nexrad": "https://upload.wikimedia.org/wikipedia/commons/9/9b/NEXRAD_radar_of_an_EF2_tornado_in_Kansas_on_March_13%2C_2024.png",
    "Tropical Tidbits GFS 500mb": "https://sites.gatech.edu/eas-mesoscale-blog/files/2023/04/Hebert_Fig4-768x531.png",
    "GOES-16 Visible": "https://cdn.star.nesdis.noaa.gov/GOES16/ABI/CONUS/GEOCOLOR/1250x750.jpg",
    "KOUN Sounding": "https://sharppy.github.io/SHARPpy/_images/gui.sharppy.png",
}

def load_image_from_url(url):
    """Load an image from URL."""
    try:
        response = requests.get(url, timeout=10, headers={'User-Agent': 'Mozilla/5.0'})
        response.raise_for_status()
        img = Image.open(BytesIO(response.content))
        # Convert to RGB if necessary
        if img.mode != 'RGB':
            img = img.convert('RGB')
        return img
    except Exception as e:
        return None, f"Error loading image from URL: {str(e)}"

@spaces.GPU(duration=90)
def analyze_weather_image(image, image_url, analysis_type, custom_prompt, analysis_mode, temperature, max_tokens, top_p):
    # Handle image input - either direct upload or URL
    if image_url and image_url.strip():
        result = load_image_from_url(image_url.strip())
        if isinstance(result, tuple):  # Error case
            return result[1]
        image = result
    elif image is None:
        return "Please upload an image or provide an image URL to analyze."
    
    # Move model to GPU
    model.cuda()
    
    # Use custom prompt if provided, otherwise use template
    prompt = custom_prompt.strip() if custom_prompt.strip() else PROMPT_TEMPLATES.get(analysis_type, PROMPT_TEMPLATES["Quick Analysis"])
    
    # Select system prompt based on mode
    system_content = SYSTEM_PROMPTS.get(analysis_mode, SYSTEM_PROMPTS["technical"])
    
    # Prepare messages
    messages = [{
        "role": "system",
        "content": system_content
    }, {
        "role": "user",
        "content": [
            {"type": "text", "text": prompt},
            {"type": "image", "image": image}
        ]
    }]
    
    # Process inputs
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text], images=[image], return_tensors="pt").to("cuda")
    
    # Generate with specified parameters
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        temperature=temperature,
        do_sample=True,
        top_p=top_p,
        repetition_penalty=1.05
    )
    
    # Decode response
    response = processor.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0]
    
    # Add metadata
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")
    source = "URL" if image_url and image_url.strip() else "Upload"
    metadata = f"\n\n---\n*Analysis completed: {timestamp} | Mode: {analysis_mode} | Type: {analysis_type} | Source: {source}*"
    
    return response + metadata

# Create Gradio interface
with gr.Blocks(title=TITLE, theme=gr.themes.Base(), css="""
    .gradio-container { font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace; }
    .markdown-text { font-size: 14px; }
    #analysis-output { font-family: 'Monaco', 'Menlo', monospace; font-size: 13px; }
    .gr-button-primary { background-color: #2563eb; }
    .gr-button-primary:hover { background-color: #1e40af; }
    .url-button { min-width: 120px; }
""") as demo:
    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Image input options
            with gr.Tabs():
                with gr.Tab("Upload Image"):
                    image_input = gr.Image(
                        label="Upload Weather Image",
                        type="pil",
                        elem_id="image-upload"
                    )
                
                with gr.Tab("Image URL"):
                    image_url_input = gr.Textbox(
                        label="Image URL",
                        placeholder="https://example.com/weather-image.jpg",
                        lines=1
                    )
                    
                    # Quick URL examples
                    gr.Markdown("**Quick Examples:**")
                    url_buttons = []
                    for name, url in EXAMPLE_URLS.items():
                        btn = gr.Button(name, size="sm", elem_classes="url-button")
                        btn.click(lambda u=url: u, outputs=image_url_input)
            
            # Analysis type selector
            analysis_type = gr.Dropdown(
                label="Analysis Type",
                choices=list(PROMPT_TEMPLATES.keys()),
                value="Quick Analysis",
                info="Select the type of analysis you need"
            )
            
            # Analysis mode selector
            analysis_mode = gr.Radio(
                label="Analysis Mode",
                choices=list(MODE_INFO.keys()),
                value="technical",
                info="Choose the style and depth of analysis"
            )
            
            # Mode description
            mode_description = gr.Markdown(value=MODE_INFO["technical"], elem_id="mode-desc")
            
            # Custom prompt option
            with gr.Accordion("Custom Prompt (Optional)", open=False):
                custom_prompt = gr.Textbox(
                    label="Enter your specific question or analysis request",
                    placeholder="E.g., 'Focus on the 500mb vorticity patterns' or 'Explain the hodograph curvature'",
                    lines=3
                )
            
            # Advanced settings
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    temperature = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.7,
                        step=0.05,
                        label="Temperature",
                        info="Lower = more focused, Higher = more varied"
                    )
                    
                    top_p = gr.Slider(
                        minimum=0.5,
                        maximum=1.0,
                        value=0.95,
                        step=0.05,
                        label="Top-p",
                        info="Nucleus sampling threshold"
                    )
                
                max_tokens = gr.Slider(
                    minimum=128,
                    maximum=1024,
                    value=512,
                    step=64,
                    label="Max Output Length",
                    info="Longer for detailed analysis"
                )
            
            # Analyze button
            analyze_btn = gr.Button("🔍 Analyze Weather", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            # Output area
            output = gr.Textbox(
                label="Analysis Results",
                lines=25,
                max_lines=30,
                show_copy_button=True,
                elem_id="analysis-output"
            )
    
    # Common weather data categories for quick access
    with gr.Accordion("📊 Quick Templates for Common Data Types", open=False):
        gr.Markdown("""
        ### Click to load analysis templates:
        """)
        with gr.Row():
            gr.Button("500mb Analysis", size="sm").click(
                lambda: "Analyze the 500mb height and wind patterns. Identify troughs, ridges, jet streaks, and vorticity.",
                outputs=custom_prompt
            )
            gr.Button("Sounding Analysis", size="sm").click(
                lambda: "Analyze this sounding for stability, CAPE, shear, LCL, LFC, and severe weather parameters.",
                outputs=custom_prompt
            )
            gr.Button("Composite Reflectivity", size="sm").click(
                lambda: "Analyze radar reflectivity patterns, storm structure, intensity, and movement.",
                outputs=custom_prompt
            )
        with gr.Row():
            gr.Button("Surface Analysis", size="sm").click(
                lambda: "Analyze surface features including fronts, pressure centers, convergence, and boundaries.",
                outputs=custom_prompt
            )
            gr.Button("Ensemble Spread", size="sm").click(
                lambda: "Analyze ensemble spread, clustering, and probabilistic information.",
                outputs=custom_prompt
            )
            gr.Button("Convective Parameters", size="sm").click(
                lambda: "Analyze CAPE, CIN, SRH, bulk shear, and composite parameters for severe potential.",
                outputs=custom_prompt
            )
    
    # Tips section
    with gr.Accordion("💡 Pro Tips for Best Results", open=False):
        gr.Markdown("""
        ### Image Guidelines:
        - **Resolution**: Higher resolution images yield better analysis
        - **Clarity**: Ensure text/contours are legible
        - **Completeness**: Include colorbars, titles, valid times
        
        ### Using URLs:
        - Supports direct links to JPG, PNG, GIF images
        - Some weather sites may block direct access
        - For best results, use official weather service URLs
        - Alternative: Save image locally and upload
        
        ### Common Weather Data Sources:
        - **SPC**: Convective outlooks, mesoanalysis
        - **College of DuPage**: NEXRAD, model data
        - **Tropical Tidbits**: Model analysis maps
        - **GOES Imagery**: Satellite data
        - **WPC**: Surface analysis, QPF
        
        ### Analysis Tips by Data Type:
        
        **Model Output (GFS, HRRR, ECMWF, NAM):**
        - Include initialization and valid times
        - Specify if you want focus on particular features
        - Ask about ensemble uncertainty if applicable
        
        **Soundings (Skew-T, Hodographs):**
        - Ensure all parameters are visible
        - Ask about specific levels or layers
        - Request shear calculations or thermodynamic analysis
        
        **Radar/Satellite:**
        - Include timestamp and location
        - Specify interest in particular features
        - Ask about storm motion or development
        """)
    
    # Update mode description when mode changes
    analysis_mode.change(
        lambda mode: MODE_INFO[mode],
        inputs=analysis_mode,
        outputs=mode_description
    )
    
    # Set up event handler
    analyze_btn.click(
        fn=analyze_weather_image,
        inputs=[image_input, image_url_input, analysis_type, custom_prompt, analysis_mode, temperature, max_tokens, top_p],
        outputs=output
    )
    
    # Clear image upload when URL is entered
    image_url_input.change(
        lambda: None,
        outputs=image_input
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()