app.py

import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import io
import base64
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime
import json
import os
import tempfile
import zipfile
import huggingface_hub
from huggingface_hub import hf_hub_download

# Import your PaveCLIP model (adjust import based on your model structure)
from paveclip_training import PaveCLIPEvaluator

# Download model from Hugging Face Hub if needed
def download_model():
    """Download model from Hugging Face Hub"""
    try:
        # Replace with your actual model repository
        model_path = hf_hub_download(
            repo_id="your-username/paveclip-model",  # Update this
            filename="paveclip_best.pt"
        )
        return model_path
    except:
        # Fallback to local path if available
        return "./paveclip_best.pt"

def download_model_from_hf():
    """Download model from separate HF model repository"""
    try:
        print("📥 Downloading PaveCLIP model...")
        model_path = hf_hub_download(
            repo_id="Blessing988/paveclip-model",  # Your model repo
            filename="paveclip_best.pt",
            cache_dir="./models"
        )
        print("✅ Model downloaded successfully!")
        return model_path
    except Exception as e:
        print(f"❌ Download failed: {e}")
        return None

class PavementAnalysisApp:
    def __init__(self, model_path):
        """Initialize the Pavement Analysis App"""
        model_path = download_model_from_hf()
        
        if model_path:
            self.evaluator = PaveCLIPEvaluator(model_path, {})
        
        # Pavement-specific class definitions
        self.distress_classes = [
            "pavement with longitudinal crack", 
            "pavement with lateral crack",
            "pavement with fatigue crack",
            "pavement with pothole",
            "road with patching",
            "pavement image showing fatigue cracks",
            "pavement image with network of cracks",
            
        ]
        
        self.material_classes = [
            "a rutting pavement image",
            "snow on a pavement",
            "asphalt road surface",
            "wet asphalt surface", 
            "wet concrete surface",
            "a severe wet concrete surface",
            "concrete road surface",
            "gravel road surface",
            "dry and smooth asphalt surface",
            "a severe dry concrete pavement surface", 
        ]
        
        self.condition_classes = [
            "smooth road surface",
            "slightly uneven road surface", 
            "severely damaged road surface",
            "well-maintained pavement",
            "deteriorated pavement"
        ]
        
        # Store embeddings for comparison
        self.image_embeddings = {}
        self.text_embeddings = {}
    
    def analyze_single_image(self, image, analysis_type="all"):
        """Analyze a single uploaded image"""
        if image is None:
            return "Please upload an image first.", {}, {}, {}
        
        # Save temporary image
        temp_path = "temp_image.jpg"
        image.save(temp_path)
        
        results = {}
        
        try:
            if analysis_type in ["distress", "all"]:
                distress_result = self.evaluator.zero_shot_classification([temp_path], self.distress_classes)
                results["distress"] = self._format_results(distress_result, self.distress_classes)
            
            if analysis_type in ["material", "all"]:
                material_result = self.evaluator.zero_shot_classification([temp_path], self.material_classes)
                results["material"] = self._format_results(material_result, self.material_classes)
            
            if analysis_type in ["condition", "all"]:
                condition_result = self.evaluator.zero_shot_classification([temp_path], self.condition_classes)
                results["condition"] = self._format_results(condition_result, self.condition_classes)
            
            # Generate summary text
            summary = self._generate_summary(results)
            
            # Clean up
            os.remove(temp_path)
            
            return summary, results.get("distress", {}), results.get("material", {}), results.get("condition", {})
            
        except Exception as e:
            os.remove(temp_path) if os.path.exists(temp_path) else None
            return f"Error analyzing image: {str(e)}", {}, {}, {}
    
    # ... (Include all other methods from the main app class)
    
    def _format_results(self, result, class_names):
        """Format classification results for display"""
        predictions = result["predictions"]
        similarities = result["similarities"]
        
        formatted = {}
        for i, class_name in enumerate(class_names):
            confidence = float(similarities[0][i])
            formatted[class_name] = confidence
        
        return formatted
    
    def _generate_summary(self, results):
        """Generate text summary of analysis"""
        summary_parts = ["🔍 **Pavement Analysis Results**\n"]
        
        for category, result in results.items():
            if result:
                best_match = max(result.items(), key=lambda x: x[1])
                category_name = category.capitalize()
                summary_parts.append(f"**{category_name}:** {best_match[0]} (confidence: {best_match[1]:.3f})")
        
        return "\n".join(summary_parts)

def create_demo():
    """Create the Gradio demo"""
    
    # Download/load model
    model_path = download_model()
    app = PavementAnalysisApp(model_path)
    
    # Create interface
    with gr.Blocks(title="🛣️ PaveCLIP: Advanced Pavement Analysis") as demo:
        
        gr.Markdown("""
        # 🛣️ PaveCLIP: Advanced Pavement Analysis Platform
        
        **Professional pavement condition assessment using state-of-the-art computer vision**
        
        Upload pavement images to get comprehensive analysis including distress detection, 
        material classification, and condition assessment.
        """)
        
        with gr.Tab("🖼️ Single Image Analysis"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="pil", label="Upload Pavement Image")
                    analysis_type = gr.Radio(
                        choices=["all", "distress", "material", "condition"],
                        value="all",
                        label="Analysis Type"
                    )
                    analyze_btn = gr.Button("🔍 Analyze Image", variant="primary")
                
                with gr.Column():
                    analysis_summary = gr.Markdown(label="Analysis Summary")
            
            with gr.Row():
                distress_output = gr.JSON(label="Distress Classification")
                material_output = gr.JSON(label="Material Classification") 
                condition_output = gr.JSON(label="Condition Assessment")
            
            analyze_btn.click(
                fn=app.analyze_single_image,
                inputs=[input_image, analysis_type],
                outputs=[analysis_summary, distress_output, material_output, condition_output]
            )
        
        # Add examples
        gr.Examples(
            examples=[
                ["examples/longitudinal-image.jpg", "distress"],
                ["examples/202202122309381-dry-asphalt-severe.jpg", "condition"],
                ["examples/202205031731377-wet-concrete-severe.jpg", "condition"],
                ["examples/202202122342019-dry-concrete-slight.jpg", "condition"]
            ],
            inputs=[input_image, analysis_type],
            outputs=[analysis_summary, distress_output, material_output, condition_output],
            fn=app.analyze_single_image,
            cache_examples=True
        )
    
    return demo

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