Upload web_app.py with huggingface_hub

web_app.py

@@ -0,0 +1,254 @@
+#!/usr/bin/env python3
+"""
+Flask Web Application — Thermal Pattern Analysis Interface.
+
+Usage:
+    python web_app.py
+    → Open http://localhost:5000
+"""
+
+import os
+import io
+import base64
+import torch
+import cv2
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import torch.nn as nn
+from pathlib import Path
+from flask import Flask, render_template, request, jsonify
+from flask_cors import CORS
+
+from src.utils.config import load_config, setup_device
+from src.preprocessing.image_processor import ThermalImageProcessor
+from src.models.anomaly_detector import ThermalPatternPipeline
+
+app = Flask(__name__)
+CORS(app)
+
+# ── Global model state ───────────────────────────────────────────────
+MODEL = None
+CLASSIFIER = None
+PROCESSOR = None
+DEVICE = None
+
+
+def load_model():
+    """Load model, classifier, and processor at startup."""
+    global MODEL, CLASSIFIER, PROCESSOR, DEVICE
+
+    config = load_config("configs/config.yaml")
+    DEVICE = setup_device(config)
+
+    MODEL = ThermalPatternPipeline.from_config(config).to(DEVICE)
+    CLASSIFIER = nn.Linear(config.model.feature_extractor.embedding_dim, 2).to(DEVICE)
+
+    ckpt_path = Path("checkpoints/best_model.pt")
+    if ckpt_path.exists():
+        ckpt = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
+        MODEL.load_state_dict(ckpt["model_state_dict"])
+        CLASSIFIER.load_state_dict(ckpt["classifier_state_dict"])
+        print(f"  ✓ Model loaded from {ckpt_path}")
+    else:
+        print(f"  ✗ No checkpoint at {ckpt_path}")
+
+    MODEL.eval()
+    CLASSIFIER.eval()
+    PROCESSOR = ThermalImageProcessor.from_config(config)
+
+
+def img_to_base64(img, cmap=None):
+    """Convert numpy image to base64-encoded PNG for HTML display."""
+    # Normalize to 0-255 uint8 if needed
+    if img.dtype == np.float32 or img.dtype == np.float64:
+        img_u8 = (np.clip(img, 0, 1) * 255).astype(np.uint8) if img.max() <= 1.0 else np.clip(img, 0, 255).astype(np.uint8)
+    else:
+        img_u8 = img.astype(np.uint8)
+
+    if cmap == 'jet':
+        # Grad-CAM heatmap
+        colored = cv2.applyColorMap(img_u8, cv2.COLORMAP_JET)
+    elif len(img_u8.shape) == 2:
+        # Grayscale → apply thermal inferno colormap
+        colored = cv2.applyColorMap(img_u8, cv2.COLORMAP_INFERNO)
+    else:
+        # Already colored (like overlay)
+        colored = cv2.cvtColor(img_u8, cv2.COLOR_RGB2BGR) if img_u8.shape[2] == 3 else img_u8
+
+    _, buf = cv2.imencode('.png', colored)
+    return base64.b64encode(buf.tobytes()).decode('utf-8')
+
+
+def compute_gradcam(input_tensor):
+    """Compute Grad-CAM heatmap."""
+    target_layer = MODEL.feature_extractor.layer4[-1].conv2
+    activations, gradients = {}, {}
+
+    def fwd_hook(m, i, o): activations["v"] = o.detach()
+    def bwd_hook(m, gi, go): gradients["v"] = go[0].detach()
+
+    fh = target_layer.register_forward_hook(fwd_hook)
+    bh = target_layer.register_full_backward_hook(bwd_hook)
+
+    try:
+        img = input_tensor.unsqueeze(0).to(DEVICE)
+        features = MODEL.feature_extractor(img)
+        MODEL.zero_grad()
+        features.max().backward()
+
+        acts = activations["v"].squeeze(0)
+        grads = gradients["v"].squeeze(0)
+        weights = grads.mean(dim=(1, 2))
+        cam = torch.relu((weights[:, None, None] * acts).sum(0))
+        cam = cam / (cam.max() + 1e-8)
+        cam = cam.cpu().numpy()
+        return cv2.resize(cam, (224, 224))
+    finally:
+        fh.remove()
+        bh.remove()
+
+
+# ── Routes ────────────────────────────────────────────────────────────
+
+@app.route("/")
+def index():
+    return render_template("index.html")
+
+
+@app.route("/analyze", methods=["POST"])
+def analyze():
+    if "file" not in request.files:
+        return jsonify({"error": "No file uploaded"}), 400
+
+    file = request.files["file"]
+    file_bytes = np.frombuffer(file.read(), np.uint8)
+    img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
+
+    if img is None:
+        return jsonify({"error": "Cannot read image"}), 400
+
+    # Grayscale
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img.copy()
+    original = gray.copy()
+
+    # Preprocessing steps
+    resized = PROCESSOR.resize(gray)
+    denoised = PROCESSOR.denoise(resized)
+    enhanced = PROCESSOR.enhance_contrast(denoised)
+    normalized = enhanced.astype(np.float32) / 255.0
+
+    # Inference
+    with torch.no_grad():
+        img_tensor = torch.from_numpy(normalized).unsqueeze(0)  # [1, H, W]
+        sequence = img_tensor.unsqueeze(0).repeat(1, 5, 1, 1).unsqueeze(2)  # [1, 5, 1, H, W]
+        sequence = sequence.to(DEVICE)
+
+        results = MODEL(sequence)
+        logits = CLASSIFIER(results["encoding"])
+        probs = torch.softmax(logits, dim=1)
+        anomaly_score = probs[0, 1].item()
+        prediction = "ABNORMAL" if anomaly_score > 0.5 else "NORMAL"
+        confidence = max(anomaly_score, 1 - anomaly_score) * 100
+
+    # Grad-CAM
+    gradcam = compute_gradcam(img_tensor)
+
+    # Create overlay
+    heatmap_colored = cv2.applyColorMap((gradcam * 255).astype(np.uint8), cv2.COLORMAP_JET)
+    base_bgr = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR)
+    overlay = cv2.addWeighted(base_bgr, 0.6, heatmap_colored, 0.4, 0)
+    overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
+
+    # Encode images
+    response = {
+        "prediction": prediction,
+        "anomaly_score": round(anomaly_score, 4),
+        "confidence": round(confidence, 1),
+        "images": {
+            "original": img_to_base64(original),
+            "resized": img_to_base64(resized),
+            "denoised": img_to_base64(denoised),
+            "enhanced": img_to_base64(enhanced),
+            "normalized": img_to_base64(normalized),
+            "gradcam": img_to_base64(gradcam, cmap='jet'),
+            "overlay": img_to_base64(overlay_rgb),
+        }
+    }
+
+    return jsonify(response)
+
+
+@app.route("/sample_images")
+def sample_images():
+    """Return list of sample images from the dataset."""
+    import glob
+    samples = glob.glob("data/raw/Power Transformers/*.jpg")[:12]
+    names = [Path(s).name for s in samples]
+    return jsonify(names)
+
+
+@app.route("/analyze_sample/<filename>")
+def analyze_sample(filename):
+    """Analyze a sample image from the dataset."""
+    path = Path("data/raw/Power Transformers") / filename
+    if not path.exists():
+        return jsonify({"error": "Sample not found"}), 404
+
+    with open(path, "rb") as f:
+        from werkzeug.datastructures import FileStorage
+        file = FileStorage(f, filename=filename)
+        # Read the file manually
+        file_bytes = np.frombuffer(f.read(), np.uint8)
+
+    img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
+    if img is None:
+        return jsonify({"error": "Cannot read image"}), 400
+
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img.copy()
+    original = gray.copy()
+    resized = PROCESSOR.resize(gray)
+    denoised = PROCESSOR.denoise(resized)
+    enhanced = PROCESSOR.enhance_contrast(denoised)
+    normalized = enhanced.astype(np.float32) / 255.0
+
+    with torch.no_grad():
+        img_tensor = torch.from_numpy(normalized).unsqueeze(0)
+        sequence = img_tensor.unsqueeze(0).repeat(1, 5, 1, 1).unsqueeze(2)
+        sequence = sequence.to(DEVICE)
+        results = MODEL(sequence)
+        logits = CLASSIFIER(results["encoding"])
+        probs = torch.softmax(logits, dim=1)
+        anomaly_score = probs[0, 1].item()
+        prediction = "ABNORMAL" if anomaly_score > 0.5 else "NORMAL"
+        confidence = max(anomaly_score, 1 - anomaly_score) * 100
+
+    gradcam = compute_gradcam(img_tensor)
+    heatmap_colored = cv2.applyColorMap((gradcam * 255).astype(np.uint8), cv2.COLORMAP_JET)
+    base_bgr = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR)
+    overlay = cv2.addWeighted(base_bgr, 0.6, heatmap_colored, 0.4, 0)
+    overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
+
+    return jsonify({
+        "prediction": prediction,
+        "anomaly_score": round(anomaly_score, 4),
+        "confidence": round(confidence, 1),
+        "images": {
+            "original": img_to_base64(original),
+            "resized": img_to_base64(resized),
+            "denoised": img_to_base64(denoised),
+            "enhanced": img_to_base64(enhanced),
+            "normalized": img_to_base64(normalized),
+            "gradcam": img_to_base64(gradcam, cmap='jet'),
+            "overlay": img_to_base64(overlay_rgb),
+        }
+    })
+
+
+if __name__ == "__main__":
+    print("Loading model...")
+    load_model()
+    port = int(os.environ.get("PORT", 5000))
+    print(f"Starting server on port {port}")
+    app.run(debug=False, host="0.0.0.0", port=port)