app.py

import os
import torch
import gradio as gr
from PIL import Image
import numpy as np
import sys
import time

# ========== 1. Import project modules and Model Configuration ==========
try:
    from stoma_clip import pmc_clip
    from stoma_clip.pmc_clip.factory import _rescan_model_configs
    from stoma_clip.training.fusion_method import convert_model_to_cls
    from stoma_clip.training.dataset.utils import encode_mlm
    print("Stoma-CLIP modules imported successfully.")
    sys.stdout.flush() # 强制刷新输出
except ImportError as e:
    print(f"FATAL: Error importing Stoma-CLIP modules: {e}")
    sys.stdout.flush()
    sys.exit(1)

# ========== 2. Model Configuration and Loading ==========
LABEL_MAP = {
    "Irritant dermatitis": 0, "Allergic contact dermatitis": 1, "Mechanical injury": 2,
    "Folliculitis": 3, "Fungal infection": 4, "Skin hyperplasia": 5, "Parastomal varices": 6,
    "Urate crystals": 7, "Cancerous metastasis": 8, "Pyoderma gangrenosum": 9, "Normal": 10
}
REVERSE_LABEL_MAP = {v: k for k, v in LABEL_MAP.items()}
NUM_CLASSES = len(LABEL_MAP)

class Args:
    def __init__(self):
        self.model = "RN50_fusion4"
        self.pretrained = "stoma_clip.pt"
        self.num_classes = NUM_CLASSES
        self.mlm = True
        self.crop_scale = 0.9
        self.context_length = 77
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")
        sys.stdout.flush() 
args = Args()

MODEL = None
PREPROCESS = None
TOKENIZER = None

def load_model():
    """Load model once in the main thread during application initialization."""
    global MODEL, PREPROCESS, TOKENIZER
    
    start_time = time.time()
    if MODEL is not None:
        return MODEL, PREPROCESS, TOKENIZER
        
    print(f"--- Starting Model Load Process at {time.strftime('%H:%M:%S')} ---")
    sys.stdout.flush() # 诊断点 1
    
    try:
        # Step 1: Create model and transforms
        print("1. Rescanning model configs and creating architecture...")
        sys.stdout.flush() # 诊断点 2
        
        _rescan_model_configs()
        model, _, preprocess = pmc_clip.create_model_and_transforms(args)
        model = convert_model_to_cls(model, num_classes=args.num_classes, fusion_method='cross_attention')
        print("2. Model architecture created. Moving to device...")
        sys.stdout.flush() # 诊断点 3
        
        # Move model architecture to GPU/CPU
        model.to(args.device).eval()
        
        # Step 2: Load weights - 必须确保 stoma_clip.pt 文件大小合理或复制完整
        print(f"3. Loading weights from {args.pretrained} to {args.device}...")
        sys.stdout.flush() # 诊断点 4 - 关键点：在执行耗时 I/O 前确保日志已输出
        
        # 强制使用 Float32 加载，然后转换为半精度，如果模型支持的话，有助于加速传输
        state_dict = torch.load(args.pretrained, map_location=args.device) 
        
        print("4. Weights file loaded. Cleaning state dict...")
        sys.stdout.flush() # 诊断点 5
        
        state_dict_clean = {k.replace("module.", "", 1): v for k, v in state_dict['state_dict'].items()}
        
        # Step 3: Apply weights
        print("5. Loading state dict into model architecture...")
        sys.stdout.flush() # 诊断点 6
        
        model.load_state_dict(state_dict_clean)
        
        # Step 4: Final setup
        tokenizer = model.tokenizer
        MODEL = model
        PREPROCESS = preprocess
        TOKENIZER = tokenizer
        
        end_time = time.time()
        print(f"✨ Stoma-CLIP Model loaded successfully! Total time: {end_time - start_time:.2f} seconds.")
        sys.stdout.flush() # 诊断点 7
        
        return MODEL, PREPROCESS, TOKENIZER
        
    except Exception as e:
        print(f"🔥 Error during model loading: {e}")
        sys.stdout.flush()
        raise RuntimeError(f"Failed to load Stoma-CLIP model: {e}")

# ========== 3. Inference Function ==========
def predict_stoma_clip(image: Image.Image, caption: str):
    # 确保在推理时调用加载模型（仅作为后备/懒加载）
    try:
        # 如果启动时加载失败，这里会再次尝试，但依赖于全局 MODEL 变量
        if MODEL is None:
             model, preprocess, tokenizer = load_model()
        else:
             model, preprocess, tokenizer = MODEL, PREPROCESS, TOKENIZER

    except RuntimeError:
        return "Model Loading Failed (See Logs)", {}
    
    # ... 原来的推理逻辑保持不变 ...
    image = image.convert("RGB")
    device = args.device
        
    # 将输入数据移动到 GPU
    image_tensor = preprocess(image).unsqueeze(0).to(device)
        
    mask_token, pad_token = '[MASK]', '[PAD]'
    vocab = [v for v in tokenizer.get_vocab().keys() if v not in tokenizer.all_special_tokens]
        
    bert_input, bert_label = encode_mlm(
        caption=caption,
        vocab=vocab,
        mask_token=mask_token,
        pad_token=pad_token,
        ratio=0.0,
        tokenizer=tokenizer,
        args=args,
    )
        
    with torch.no_grad():
        inputs = {"images": image_tensor, "bert_input": bert_input, "bert_label": bert_label}
        outputs = model(inputs)
        # 将结果移回 CPU 进行 numpy 转换
        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
        predicted_class_idx = torch.argmax(outputs, dim=1).item()
            
    predicted_class_name = REVERSE_LABEL_MAP.get(predicted_class_idx, "Unknown")
    probability_distribution = {REVERSE_LABEL_MAP[i]: float(p) for i, p in enumerate(probs)}
    return predicted_class_name, probability_distribution

# ========== 4. Gradio Interface Setup ==========
image_input = gr.Image(type="pil", label="上传造口图片")
caption_input = gr.Textbox(label="输入造口描述文本 (例如: Exudate, epidermal breakdown, ...)")
predicted_label_output = gr.Textbox(label="预测类别")
prob_output = gr.Label(label="类别概率分布")

# Find example path for Gradio demo (No Change)
try:
    example_path_1 = "demo/Irritant_dermatitis.jpg"
    example_path_2 = "demo/Folliculitis.jpg"
    examples_list = []
    if os.path.exists(example_path_1):
        examples_list.append(
            [example_path_1, "Exudate, epidermal breakdown, irregular erythema, pain, confined to contact areas"])
    elif os.path.exists(example_path_2):
        examples_list.append([example_path_2, "Erythema, papules, pustules confined to hair follicles"])
except Exception:
    examples_list = []

# 忽略 Gradio 的 allow_flagging 警告
iface = gr.Interface(
    fn=predict_stoma_clip,
    inputs=[image_input, caption_input],
    outputs=[predicted_label_output, prob_output],
    title="🧪 Stoma-CLIP 分类 API 原型 (Gradio)",
    description="请上传造口图片并输入临床描述，模型将预测最可能的皮肤并发症类别。",
    examples=examples_list,
    allow_flagging="never"
)

if __name__ == "__main__":
    # --- 关键修复：强制在 Gradio launch 之前加载模型，将 I/O 阻塞移到启动阶段 ---
    print("Pre-loading model before Gradio launch to prevent runtime timeout...")
    sys.stdout.flush()
    
    load_model() 
    
    print("Model loaded. Launching Gradio interface...")
    sys.stdout.flush()

    # 启动 Gradio
    iface.launch(server_name="0.0.0.0", server_port=7860)