app.py

import os
import sys
import tempfile
import cv2
import dotenv
import gradio as gr

# Import the image slider component
try:
    from gradio_imageslider import ImageSlider
except ImportError:
    raise ImportError("Please install the slider component by running: pip install gradio-imageslider")

# Load environment variables (e.g. GROQ_API_KEY)
dotenv.load_dotenv()

# Add app.py directory and backend/ directory to system path so imports work in all layouts
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)
sys.path.append(os.path.join(current_dir, "backend"))

# Fallback imports in case backend files aren't found in your environment
try:
    from diff_ai import compare_images_ai
    from diff import compare_images
    from describe import describe_change
except ImportError:
    print("Warning: Backend modules not found. Using mock functions.")
    def compare_images_ai(*args, **kwargs): return {"ssim_score": 0.95, "severity": "low"}
    def compare_images(*args, **kwargs): return {"ssim_score": 0.95, "severity": "low"}
    def describe_change(*args, **kwargs): return "Mock description of changes."

def gradio_compare(img1, img2):
    if img1 is None or img2 is None:
        return None, None, None, None, "Please upload both images.", {}

    with tempfile.TemporaryDirectory() as tmpdir:
        path1 = os.path.join(tmpdir, "img1.jpg")
        path2 = os.path.join(tmpdir, "img2.jpg")
        output_path = os.path.join(tmpdir, "result.png")
        output_box_path = os.path.join(tmpdir, "result_box.png")
        heatmap_path = os.path.join(tmpdir, "heatmap.png")

        # Convert RGB inputs (Gradio) to BGR (OpenCV)
        img1_bgr = cv2.cvtColor(img1, cv2.COLOR_RGB2BGR)
        img2_bgr = cv2.cvtColor(img2, cv2.COLOR_RGB2BGR)

        # Save images
        cv2.imwrite(path1, img1_bgr)
        cv2.imwrite(path2, img2_bgr)

        # Run comparison pipeline
        pipeline_used = "AI Pipeline"
        try:
            result_data = compare_images_ai(path1, path2, output_path, heatmap_path, output_box_path=output_box_path)
        except Exception as e:
            print(f"[Gradio] AI pipeline failed ({e}); falling back to classical")
            result_data = compare_images(path1, path2, output_path, heatmap_path)
            if os.path.exists(output_path):
                import shutil
                shutil.copy(output_path, output_box_path)
            result_data["pipeline"] = "classical"
            pipeline_used = "Classical Pipeline"

        # Generate description
        try:
            description = describe_change(path1, path2, result_data)
        except Exception as e:
            description = f"Error generating description: {e}"

        # Load output images (convert BGR back to RGB for Gradio display)
        result_img = cv2.cvtColor(cv2.imread(output_path), cv2.COLOR_BGR2RGB) if os.path.exists(output_path) else None
        box_img = cv2.cvtColor(cv2.imread(output_box_path), cv2.COLOR_BGR2RGB) if os.path.exists(output_box_path) else None
        heatmap_img = cv2.cvtColor(cv2.imread(heatmap_path), cv2.COLOR_BGR2RGB) if os.path.exists(heatmap_path) else None

        # Format metrics dictionary
        metrics = {
            "Pipeline Used": pipeline_used,
            "SSIM Score (Similarity)": f"{result_data.get('ssim_score', 0):.4f}" if isinstance(result_data.get('ssim_score'), (int, float)) else "N/A",
            "LPIPS Distance (Perceptual)": f"{result_data.get('lpips_max', 0):.4f}" if isinstance(result_data.get('lpips_max'), (int, float)) else "N/A",
            "Detected Object Changes": result_data.get("object_changes", "0"),
            "Detected Crack Count": result_data.get("crack_count", "0"),
            "Change Severity": str(result_data.get("severity", "none")).upper(),
            "Alignment Failed": "Yes" if result_data.get("alignment_failed") else "No"
        }

        # Return slider tuple (img1, img2), plus the other images and texts
        return (img1, img2), result_img, box_img, heatmap_img, description, metrics


# --- UI Theme and Layout Design ---
# Switched to Base theme for a more minimal starting point
theme = gr.themes.Base(
    primary_hue="indigo",
    secondary_hue="slate",
    neutral_hue="slate"
)

css_styling = """
/* Container limit */
.gradio-container { max-width: 1600px !important; }

/* Force Tailwind text-sm (14px, line-height 20px) on text elements */
.gradio-container, 
.gradio-container .prose *, 
.gradio-container button, 
.gradio-container span, 
.gradio-container label, 
.gradio-container input, 
.gradio-container textarea, 
.gradio-container p {
    font-size: 0.875rem !important; 
    line-height: 1.25rem !important;
}

/* Ensure headings are also text-sm but bolder to maintain hierarchy */
.gradio-container h1, .gradio-container h2, .gradio-container h3, .gradio-container h4 {
    font-size: 0.875rem !important;
    font-weight: 600 !important;
    margin-bottom: 0.5rem !important;
}

/* Sidebar & Dark Panel styling */
.dark-panel { 
    padding: 20px;
    border-radius: 12px; 
    background: rgba(0, 0, 0, 0.2) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
}

/* Minimal Process Button */
.process-btn { 
    margin-top: 20px; 
    font-weight: 600 !important; 
    border-radius: 8px !important; 
    background: rgba(255, 255, 255, 0.1) !important;
    border: 1px solid rgba(255, 255, 255, 0.2) !important;
}
.process-btn:hover {
    background: rgba(255, 255, 255, 0.15) !important;
}

/* Remove default Gradio block borders for an ultra-minimal look */
.gr-block { border: none !important; }
"""

with gr.Blocks(theme=theme, css=css_styling) as demo:
    
    with gr.Row():
        # --- LEFT COLUMN: Input Controls (Scale 1 for fixed-like width) ---
        with gr.Column(scale=1, min_width=280, elem_classes=["dark-panel"]):
            gr.Markdown("### Input Images")
            
            img1_input = gr.Image(label="BEFORE IMAGE", type="numpy", height=250)
            img2_input = gr.Image(label="AFTER IMAGE", type="numpy", height=250)
                
            submit_btn = gr.Button("Process Analysis", variant="primary", elem_classes=["process-btn"])
            
        # --- RIGHT COLUMN: Unified Comparison Canvas (Scale 5 to consume all remaining space) ---
        with gr.Column(scale=5, elem_classes=["dark-panel"]):
            gr.Markdown("### Outputs")
                
            # Image Slider component
            slider_output = ImageSlider(label="Before / After", type="numpy", show_label=False)
            
            # Visuals Tabs
            with gr.Tabs():
                with gr.TabItem("Segment Overlays"):
                    result_img_output = gr.Image(label="Segment Overlay Highlight")
                with gr.TabItem("Bounding Boxes"):
                    result_box_output = gr.Image(label="Bounding Box Detections")
                with gr.TabItem("Difference Heatmap"):
                    heatmap_img_output = gr.Image(label="Difference Heatmap")
                    
            # --- LOWER SECTION: AI Analysis & Metrics (Single Column) ---
            gr.Markdown("### AI Analysis & Metrics")
            desc_output = gr.Textbox(label="Explanation of Changes", lines=4)
            metrics_output = gr.JSON(label="Quantitative Comparison")

    # Connect the backend function
    submit_btn.click(
        fn=gradio_compare,
        inputs=[img1_input, img2_input],
        outputs=[
            slider_output, 
            result_img_output, 
            result_box_output, 
            heatmap_img_output, 
            desc_output, 
            metrics_output
        ]
    )

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