Upload app.py

app.py

@@ -0,0 +1,88 @@
+import json
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from torchvision.models import efficientnet_b0
+from PIL import Image
+import gradio as gr
+
+
+# ---------- CONFIG ----------
+CKPT_PATH = "best_effnet_twohead.pt"
+LABELS_PATH = "labels.json"
+IMG_SIZE = 224
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# ----------------------------
+
+
+# load labels
+with open(LABELS_PATH, "r") as f:
+    labels = json.load(f)
+
+SPECIES = labels["species"]
+STATE = labels["state"]
+
+
+# model (must match training)
+class EffNetTwoHead(nn.Module):
+    def __init__(self, num_species, num_states):
+        super().__init__()
+        base = efficientnet_b0(weights=None)
+        self.features = base.features
+        self.avgpool = base.avgpool
+        c = base.classifier[1].in_features
+        self.head_species = nn.Linear(c, num_species)
+        self.head_state = nn.Linear(c, num_states)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        return self.head_species(x), self.head_state(x)
+
+
+# load model
+ckpt = torch.load(CKPT_PATH, map_location="cpu")
+model = EffNetTwoHead(len(SPECIES), len(STATE))
+model.load_state_dict(ckpt["model"])
+model.to(DEVICE).eval()
+
+
+# preprocessing (same as training)
+tfm = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225],
+    ),
+])
+
+
+@torch.no_grad()
+def predict(image: Image.Image):
+    if image is None:
+        return "No image", "No image"
+
+    image = image.convert("RGB")
+    x = tfm(image).unsqueeze(0).to(DEVICE)
+
+    log_sp, log_st = model(x)
+    sp_id = int(log_sp.argmax(dim=1))
+    st_id = int(log_st.argmax(dim=1))
+
+    return SPECIES[sp_id], STATE[st_id]
+
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Textbox(label="Predicted species"),
+        gr.Textbox(label="Predicted state"),
+    ],
+    title="EfficientNet Two-Head Layer Trap Nest (LTN) Classifier",
+)
+
+if __name__ == "__main__":
+    demo.launch()