main.py

import os
import sys
import json
import random
import shutil
import hashlib
import uuid
from typing import List
import base64
from io import BytesIO
import time
import threading
import numpy as np
import torch
import torch.nn as nn
from PIL import Image, ImageOps
from matplotlib import cm
import requests   # <--- เพิ่มเพื่อโหลดจาก HuggingFace

import cv2
from fastapi import FastAPI, File, UploadFile, Form, Request, Depends
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles

# ============ เพิ่มระบบดาวน์โหลดโมเดลจาก HuggingFace ============
HF_MODEL_URL = "https://huggingface.co/qqqqqqat/densenet_wangchan/resolve/main/best_fusion_densenet.pth"
LOCAL_MODEL_PATH = "models/densenet/best_fusion_densenet.pth"

def download_model_if_needed():
    if not os.path.exists(LOCAL_MODEL_PATH):
        print("📥 Downloading model from HuggingFace...")
        os.makedirs(os.path.dirname(LOCAL_MODEL_PATH), exist_ok=True)
        response = requests.get(HF_MODEL_URL)
        with open(LOCAL_MODEL_PATH, "wb") as f:
            f.write(response.content)
        print("✅ Model downloaded from HuggingFace!")
# ===================================================================

sys.path.append(os.path.abspath(os.path.dirname(__file__)))
from models.densenet.preprocess.preprocessingwangchan import get_tokenizer, get_transforms
from models.densenet.train_densenet_only import DenseNet121Classifier
from models.densenet.train_text_only import TextClassifier

torch.manual_seed(42); np.random.seed(42); random.seed(42)

FUSION_LABELMAP_PATH = "models/densenet/label_map_fusion_densenet.json"

with open(FUSION_LABELMAP_PATH, "r", encoding="utf-8") as f:
    label_map = json.load(f)

class_names = [label for label, _ in sorted(label_map.items(), key=lambda x: x[1])]
NUM_CLASSES = len(class_names)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🧠 Using device: {device}")

# ====================== Model Fusion Class ==========================
class FusionDenseNetText(nn.Module):
    def __init__(self, num_classes, dropout=0.3):
        super().__init__()
        self.image_model = DenseNet121Classifier(num_classes=num_classes)
        self.text_model  = TextClassifier(num_classes=num_classes)

        self.fusion = nn.Sequential(
            nn.Linear(num_classes * 2, 128), nn.ReLU(),
            nn.Dropout(dropout), nn.Linear(128, num_classes)
        )

    def forward(self, image, input_ids, attention_mask):
        logits_img = self.image_model(image)
        logits_txt = self.text_model(input_ids, attention_mask)
        fused_in   = torch.cat([logits_img, logits_txt], dim=1)
        fused_out  = self.fusion(fused_in)
        return fused_out, logits_img, logits_txt

# ===================== Load Model ============================
print("🔄 Loading AI model...")

# โหลดไฟล์โมเดลจาก HuggingFace ถ้ายังไม่มี
download_model_if_needed()

fusion_model = FusionDenseNetText(num_classes=NUM_CLASSES).to(device)
fusion_model.load_state_dict(torch.load(LOCAL_MODEL_PATH, map_location=device))
fusion_model.eval()

print("✅ AI Model loaded successfully!")
# =============================================================

tokenizer = get_tokenizer()
transform = get_transforms((224, 224))

def _find_last_conv2d(mod: torch.nn.Module):
    last = None
    for m in mod.modules():
        if isinstance(m, torch.nn.Conv2d): last = m
    return last

def compute_gradcam_overlay(img_pil, image_tensor, target_class_idx):
    img_branch = fusion_model.image_model
    target_layer = _find_last_conv2d(img_branch)
    if target_layer is None:
        return None

    activations, gradients = [], []

    def fwd_hook(_m, _i, o): activations.append(o)
    def bwd_hook(_m, gin, gout): gradients.append(gout[0])

    h1 = target_layer.register_forward_hook(fwd_hook)
    h2 = target_layer.register_full_backward_hook(bwd_hook)

    try:
        img_branch.zero_grad()
        logits_img = img_branch(image_tensor)
        score = logits_img[0, target_class_idx]
        score.backward()

        act = activations[-1].detach()[0]
        grad = gradients[-1].detach()[0]
        weights = torch.mean(grad, dim=(1, 2))

        cam = torch.relu(torch.sum(weights[:, None, None] * act, dim=0))
        cam -= cam.min()
        cam /= (cam.max() + 1e-8)

        cam_img = Image.fromarray((cam.cpu().numpy() * 255).astype(np.uint8)).resize(img_pil.size, Image.BILINEAR)

        heatmap = cm.get_cmap("jet")(cam_img)[:, :, :3]

        img_np = np.asarray(img_pil.convert("RGB")).astype(np.float32) / 255.0

        overlay = (0.6 * img_np + 0.4 * heatmap)
        return np.clip(overlay * 255, 0, 255).astype(np.uint8)

    finally:
        h1.remove()
        h2.remove()
        img_branch.zero_grad()


# ==================== FastAPI Server ==========================
app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")
os.makedirs("uploads", exist_ok=True)

# Cache system
EXPIRATION_MINUTES = 10 
results_cache = {}       
cache_lock = threading.Lock() 

def cleanup_expired_cache():
    while True:
        with cache_lock:
            expired_keys = []
            current_time = time.time()
            for key, value in results_cache.items():
                if current_time - value["created_at"] > EXPIRATION_MINUTES * 60:
                    expired_keys.append(key)

            for key in expired_keys:
                del results_cache[key]
                print(f"🧹 Cache expired and removed for key: {key}")

        time.sleep(60)

@app.on_event("startup")
async def startup_event():
    cleanup_thread = threading.Thread(target=cleanup_expired_cache, daemon=True)
    cleanup_thread.start()
    print("🗑️ Cache cleanup task started.")

SYMPTOM_MAP = {
    "noSymptoms": "ไม่มีอาการ", "drinkAlcohol": "ดื่มเหล้า", "smoking": "สูบบุหรี่",
    "chewBetelNut": "เคี้ยวหมาก", "eatSpicyFood": "กินเผ็ดแสบ", "wipeOff": "เช็ดออกได้",
    "alwaysHurts": "เจ็บเมื่อโดนแผล"
}

def process_with_ai_model(image_path: str, prompt_text: str):
    try:
        image_pil = Image.open(image_path)
        image_pil = ImageOps.exif_transpose(image_pil)
        image_pil = image_pil.convert("RGB")

        image_tensor = transform(image_pil).unsqueeze(0).to(device)
        enc = tokenizer(prompt_text, return_tensors="pt", padding="max_length",
                        truncation=True, max_length=128)

        ids, mask = enc["input_ids"].to(device), enc["attention_mask"].to(device)

        with torch.no_grad():
            fused_logits, _, _ = fusion_model(image_tensor, ids, mask)
            probs_fused = torch.softmax(fused_logits, dim=1)[0].cpu().numpy()

        pred_idx   = int(np.argmax(probs_fused))
        pred_label = class_names[pred_idx]
        confidence = float(probs_fused[pred_idx]) * 100

        gradcam_overlay_np = compute_gradcam_overlay(image_pil, image_tensor, pred_idx)

        def image_to_base64(img):
            buffered = BytesIO()
            img.save(buffered, format="JPEG")
            return base64.b64encode(buffered.getvalue()).decode('utf-8')

        original_b64 = image_to_base64(image_pil)

        if gradcam_overlay_np is not None:
            gradcam_pil = Image.fromarray(gradcam_overlay_np)
            gradcam_b64 = image_to_base64(gradcam_pil)
        else:
            gradcam_b64 = original_b64

        return original_b64, gradcam_b64, pred_label, f"{confidence:.2f}"

    except Exception as e:
        print(f"❌ Error during AI processing: {e}")
        return None, None, "Error", "0.00"


@app.get("/", response_class=RedirectResponse)
async def root():
    return RedirectResponse(url="/detect")

@app.get("/detect", response_class=HTMLResponse)
async def show_upload_form(request: Request):
    return templates.TemplateResponse("detect.html", {"request": request})

@app.post("/uploaded")
async def handle_upload(
    request: Request,
    file: UploadFile = File(...),
    checkboxes: List[str] = Form([]),
    symptom_text: str = Form("")
):
    temp_filepath = os.path.join("uploads", f"{uuid.uuid4()}_{file.filename}")

    with open(temp_filepath, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    selected_symptoms_thai = {SYMPTOM_MAP.get(cb) for cb in checkboxes if SYMPTOM_MAP.get(cb)}

    final_prompt_parts = []

    if "ไม่มีอาการ" in selected_symptoms_thai:
        symptoms_group   = {"เจ็บเมื่อโดนแผล", "กินเผ็ดแสบ"}
        lifestyles_group = {"ดื่มเหล้า", "สูบบุหรี่", "เคี้ยวหมาก"}
        patterns_group   = {"เช็ดออกได้"}
        special_group    = {"ไม่มีอาการ"}

        final_selected = (selected_symptoms_thai - symptoms_group) | \
                         (selected_symptoms_thai & (lifestyles_group | patterns_group | special_group))

        final_prompt_parts.append(" ".join(sorted(list(final_selected))))
    elif selected_symptoms_thai:
        final_prompt_parts.append(" ".join(sorted(list(selected_symptoms_thai))))

    if symptom_text and symptom_text.strip():
        final_prompt_parts.append(symptom_text.strip())

    final_prompt = "; ".join(final_prompt_parts) if final_prompt_parts else "ไม่มีอาการ"

    image_b64, gradcam_b64, name_out, eva_output = process_with_ai_model(
        image_path=temp_filepath, prompt_text=final_prompt
    )

    os.remove(temp_filepath)

    result_id = str(uuid.uuid4())
    result_data = {
        "image_b64_data": image_b64,
        "gradcam_b64_data": gradcam_b64,
        "name_out": name_out,
        "eva_output": eva_output,
    }

    with cache_lock:
        results_cache[result_id] = {
            "data": result_data,
            "created_at": time.time()
        }

    results_url = request.url_for('show_results', result_id=result_id)
    return RedirectResponse(url=results_url, status_code=303)


@app.get("/results/{result_id}", response_class=HTMLResponse)
async def show_results(request: Request, result_id: str):
    with cache_lock:
        cached_item = results_cache.get(result_id)

    if not cached_item or (time.time() - cached_item["created_at"] > EXPIRATION_MINUTES * 60):
        if cached_item:
            with cache_lock:
                del results_cache[result_id]
        return RedirectResponse(url="/detect")

    context = {"request": request, **cached_item["data"]}
    return templates.TemplateResponse("detect.html", context)


# =============== รองรับ Render / Railway / VPS ================
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 8000))
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=port)