Update app.py

app.py

@@ -8,21 +8,48 @@ from sklearn.linear_model import LogisticRegression
 
 DEVICE = torch.device("cpu")
 
+# ---------------- LOAD ENCODER ----------------
 encoder = load_encoder("encoder_resnet18_simclr.pth")
 
+# ---------------- LOAD LINEAR PROBE ----------------
 data = np.load("linear_probe_cifar10.npz", allow_pickle=True)
-clf = data["clf"].item()
+keys = list(data.files)
 
+print("Loaded NPZ keys:", keys)
+
+clf = None
+
+# Case 1: saved full sklearn model
+for k in keys:
+    if isinstance(data[k].item(), LogisticRegression):
+        clf = data[k].item()
+
+# Case 2: saved coef / intercept
+if clf is None and "coef" in keys and "intercept" in keys:
+    clf = LogisticRegression()
+    clf.coef_ = data["coef"]
+    clf.intercept_ = data["intercept"]
+    clf.classes_ = np.arange(10)
+
+# If still no classifier → FAIL gracefully
+if clf is None:
+    raise ValueError(f"❌ Could not find classifier inside NPZ. Found keys: {keys}")
+
+print("Classifier Loaded Successfully")
+
+# ---------------- CLASSES ----------------
 CLASSES = [
     "airplane","automobile","bird","cat","deer",
     "dog","frog","horse","ship","truck"
 ]
 
+# ---------------- TRANSFORM ----------------
 transform = transforms.Compose([
     transforms.Resize((32, 32)),
     transforms.ToTensor()
 ])
 
+# ---------------- PREDICT ----------------
 def predict(image):
     image = Image.fromarray(image).convert("RGB")
     x = transform(image).unsqueeze(0).to(DEVICE)
@@ -30,11 +57,12 @@ def predict(image):
     with torch.no_grad():
         emb = encoder(x).cpu().numpy()
 
-    pred = clf.predict(emb)[0]
     probs = clf.predict_proba(emb)[0]
+    pred = np.argmax(probs)
 
     return {CLASSES[i]: float(probs[i]) for i in range(10)}
 
+# ---------------- UI ----------------
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy"),