Upload 4 files

app.py

@@ -0,0 +1,20 @@
+import gradio as gr
+from tongue_model import TongueModelWrapper 
+
+wrapper = TongueModelWrapper(model_path="best_model.pt")
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🏥 中醫舌象自動診斷系統")
+    
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(label="上傳舌象照片")
+            btn = gr.Button("🚀 開始分析", variant="primary")
+        
+        with gr.Column():
+            output_label = gr.Label(label="預測結果機率")
+
+    btn.click(fn=wrapper.predict, inputs=input_img, outputs=output_label)
+
+if __name__ == "__main__":
+    demo.launch()

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3326eba752731c60c77defabf73e8e21bb6a23ced04285462e8c67d5ec4d71df
+size 46996455

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+gradio
+opencv-python-headless
+numpy
+Pillow

tongue_model.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import cv2
+import numpy as np
+from PIL import Image
+from torchvision import models, transforms
+
+# --- 模型架構定義 ---
+class CBAM(nn.Module):
+    def __init__(self, channels, reduction=16):
+        super(CBAM, self).__init__()
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(channels, channels // reduction, 1, bias=False),
+            nn.ReLU(),
+            nn.Conv2d(channels // reduction, channels, 1, bias=False)
+        )
+        self.sa = nn.Sequential(
+            nn.Conv2d(2, 1, kernel_size=7, padding=3, bias=False),
+            nn.Sigmoid()
+        )
+        self.ca_sigmoid = nn.Sigmoid()
+    def forward(self, x):
+        x = x * self.ca_sigmoid(self.ca(x))
+        avg_out = torch.mean(x, dim=1, keepdim=True); max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = x * self.sa(torch.cat([avg_out, max_out], dim=1))
+        return x
+
+class ArcMarginProduct(nn.Module):
+    def __init__(self, in_features, out_features, s=35.0, m=0.50):
+        super(ArcMarginProduct, self).__init__()
+        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
+        nn.init.xavier_uniform_(self.weight)
+        self.s = s
+    def forward(self, input):
+        cosine = F.linear(F.normalize(input), F.normalize(self.weight))
+        return cosine * self.s
+
+class TongueArcResNet(nn.Module):
+    def __init__(self, num_classes=3):
+        super().__init__()
+        self.backbone = models.resnet18(weights=None)
+        self.features = nn.Sequential(*list(self.backbone.children())[:-2])
+        self.attention = CBAM(512)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.arcface = ArcMarginProduct(512, num_classes, s=35)
+    def forward(self, x):
+        x = self.features(x)
+        x = self.attention(x)
+        features = self.avgpool(x).flatten(1)
+        return self.arcface(features)
+
+# --- 2. 定義預處理與推論類別 ---
+class TongueModelWrapper:
+    def __init__(self, model_path, num_classes=2):
+        self.device = torch.device("cpu")
+        self.model = TongueArcResNet(num_classes=num_classes)
+        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
+        self.model.eval()
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+    def preprocess(self, img_array):
+        img_gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+        img_gray = cv2.resize(img_gray, (512, 512))
+        
+        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+        ch_clahe = clahe.apply(img_gray)
+        
+        ch_lap = np.absolute(cv2.Laplacian(img_gray, cv2.CV_64F, ksize=3))
+        ch_lap = cv2.normalize(ch_lap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
+        
+        combined = np.stack([img_gray, ch_clahe, ch_lap], axis=-1)
+        return Image.fromarray(combined)
+
+    def predict(self, img_array):
+        if img_array is None: return None
+        
+        processed_img = self.preprocess(img_array)
+        input_tensor = self.transform(processed_img).unsqueeze(0).to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.model(input_tensor)
+            probs = torch.softmax(outputs, dim=1).numpy()[0]
+        
+        return {"NHC(健康人)":float(probs[0]), "DES (一般乾眼)": float(probs[1]), "SJS (乾燥症)": float(probs[2])}