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 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

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
import requests
import uvicorn

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.")


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


FUSION_LABELMAP_PATH = "models/densenet/label_map_fusion_densenet.json"
FUSION_WEIGHTS_PATH  = "models/densenet/best_fusion_densenet.pth"
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}")
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
print("🔄 Loading AI model...")
fusion_model = FusionDenseNetText(num_classes=NUM_CLASSES).to(device)
# ดาวน์โหลดก่อน
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!")
fusion_model.eval()
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)
        cam_np = np.asarray(cam_img).astype(np.float32) / 255.0
        heatmap = cm.get_cmap("jet")(cam_np)[:, :, :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()


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

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

def cleanup_expired_cache():
    """
    ฟังก์ชันนี้จะทำงานใน Background Thread เพื่อตรวจสอบและลบ Cache ที่หมดอายุ
    """
    while True:
        with cache_lock: # ล็อคเพื่อความปลอดภัยในการเข้าถึง cache
            # สร้าง list ของ key ที่จะลบ เพื่อไม่ให้แก้ไข dict ขณะวน loop
            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)
            
            # ลบ key ที่หมดอายุ
            for key in expired_keys:
                del results_cache[key]
                print(f"🧹 Cache expired and removed for key: {key}")
        
        time.sleep(60) # ตรวจสอบทุกๆ 60 วินาที

@app.on_event("startup")
async def startup_event():
    """
    เริ่ม Background Thread สำหรับทำความสะอาด Cache เมื่อแอปเริ่มทำงาน
    """
    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)
    final_prompt_parts = []
    selected_symptoms_thai = {SYMPTOM_MAP.get(cb) for cb in checkboxes if SYMPTOM_MAP.get(cb)}
    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)

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 8000))   # ใช้ PORT ของ hosting ถ้ามี
    uvicorn.run(app, host="0.0.0.0", port=port)