Upload app.py

app.py

@@ -11,8 +11,6 @@ import matplotlib.patches as mpatches
 import io
 
 # ── Config (from your training code) ──────────────────────────────
-MODEL_REPO = "Thomaslam1202/Cityscapes_Segmentation_Model"
-MODEL_FILE = "best.pth"
 IMG_SIZE   = (512, 512)       # your cfg.img_size
 DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -55,27 +53,20 @@ def mask_to_colour(mask: np.ndarray) -> np.ndarray:
 
 # ── Load model once at startup ─────────────────────────────────────
 def load_model():
-    print(f"Loading on: {DEVICE}")
     model = SegformerForSemanticSegmentation.from_pretrained(
-        "nvidia/mit-b2",                   # same backbone as your training
-        num_labels=19,                     # cfg.num_class
+        "nvidia/mit-b2",
+        num_labels=19,
         id2label=id_to_label,
         label2id=label_to_id,
-        ignore_mismatched_sizes=True,      # same flag as your training
+        ignore_mismatched_sizes=True,
     )
-    model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
-    ckpt = torch.load(model_path, map_location=DEVICE)
-
-    # ← exact same _orig_mod stripping logic from your resume block
+    ckpt = torch.load("best.pth", map_location=DEVICE)
     state_dict = ckpt["state_dict"]
     if any(k.startswith("_orig_mod.") for k in state_dict.keys()):
         state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}
-
     model.load_state_dict(state_dict)
     return model.to(DEVICE).eval()
 
-model = load_model()
-
 # ── Inference (mirrors your predict() function) ────────────────────
 def run_inference(pil_image: Image.Image):
     img_tensor = preprocess(pil_image.convert("RGB")).unsqueeze(0).to(DEVICE)