Update app.py

app.py

@@ -1,33 +1,39 @@
 import gradio as gr
 import torch
 import timm
+import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
-
-# モデル読み込み
 from huggingface_hub import hf_hub_download
-import torch
-import timm
 
+# デバイス設定
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# モデル構築（学習時と同じ構造を再現）
+def build_model():
+    backbone = timm.create_model('convnext_small', pretrained=False, num_classes=0, global_pool='avg')
+    model = nn.Sequential(
+        backbone,
+        nn.Linear(backbone.num_features, 128),
+        nn.ReLU(),
+        nn.Dropout(0.4),
+        nn.Linear(128, 1)
+    )
+    return model
+
+# モデルの重みを Model Hub からダウンロード
 model_path = hf_hub_download(
-    repo_id="lmedz/ips-model-weights",  # ←作ったModel Hubのリポジトリ名
+    repo_id="ryumiyake/ips-model-weights",  # あなたのModel Hubリポジトリ
     filename="model.pth"
 )
 
-
-# 例：保存時が nn.Sequential(model) だった場合
-model_base = timm.create_model('convnext_small', pretrained=False, num_classes=2)
-model = torch.nn.Sequential(model_base)  # ← wrapperを追加
-
+# モデル初期化・重み読み込み
+model = build_model()
 state_dict = torch.load(model_path, map_location=device)
 model.load_state_dict(state_dict)
 model.to(device)
 model.eval()
 
-
-
 # 前処理
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -37,17 +43,17 @@ transform = transforms.Compose([
         std=[0.229, 0.224, 0.225])
 ])
 
-class_names = ['Low CPM Score', 'High CPM Score']  # 適宜修正
-
+# 推論関数
+THRESHOLD = 0.5  # 必要に応じて変更可能
 def predict(img):
     img_tensor = transform(img).unsqueeze(0).to(device)
     with torch.no_grad():
-        logits = model(img_tensor)
-        probs = torch.softmax(logits, dim=1).cpu().numpy().flatten()
-    result = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
-    return img, result  # ← 画像とスコアを返す
-
+        logit = model(img_tensor)
+        prob = torch.sigmoid(logit).item()  # 出力: 0〜1の確率
+    label = "High CPM Score" if prob >= THRESHOLD else "Low CPM Score"
+    return img, {label: prob}
 
+# Gradioインターフェース
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Input Image"),
@@ -56,9 +62,9 @@ demo = gr.Interface(
         gr.Label(num_top_classes=2, label="Prediction")
     ],
     title="iPS Cell Quality Classifier",
-    description="Upload a microscopy image. The image and predicted cell quality will be displayed."
+    description="Upload a microscopy image to classify cell quality based on CPM score."
 )
 
-
 if __name__ == "__main__":
     demo.launch()
+