Update app.py

app.py

@@ -46,22 +46,9 @@ dinov2 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14_reg')
 print("Loading custom pathology checkpoint...")
 #ours = torch.load("/data/linears/teacher_checkpoint.pth")
 
-ours = torch.load("/data/linears/teacher_checkpoint_load.pt")
-checkpoint = ours#["teacher"]
-checkpoint_new = {}
+checkpoint = torch.load("./teacher_checkpoint_load.pt")
 
-# Remove dino and ibot keys
-#for key in list(checkpoint.keys()):
-#    if "dino" in str(key) or "ibot" in str(key):
-#        checkpoint.pop(key, None)
 
-# Align keys with dinov2 state dict
-#for key, keyb in zip(checkpoint.keys(), dinov2.state_dict().keys()):
-#    checkpoint_new[keyb] = checkpoint[key]
-
-#checkpoint = checkpoint_new
-
-# Update pos_embed
 
 new_shape = checkpoint["pos_embed"]
 dinov2.pos_embed = torch.nn.parameter.Parameter(new_shape)
@@ -102,16 +89,16 @@ def setup_linear_crc(path):
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dinov2 = dinov2.to(device)
 
-breakhis_path = "/data/linears/logs/dino_vits16/offline/breakhis/20251030-190409559498_5b47c293/run_0/checkpoints/best.ckpt"
+breakhis_path = "./breakhis_best.ckpt"
 breakhis_linear = setup_linear(breakhis_path).to(device)
 
-gleason_path = "/data/linears/logs/dino_vits16/offline/gleason_arvaniti/20251110-164046988851_35daf081/run_4/checkpoints/best.ckpt"
+gleason_path = "./gleason_best.ckpt"
 gleason_linear = setup_linear(gleason_path).to(device)
 
-bach_path = "/data/linears/logs/dino_vits16/offline/bach/20251110-164046453320_0b82d41d/run_4/checkpoints/best.ckpt"
+bach_path = "./bach_best.ckpt"
 bach_linear = setup_linear(bach_path).to(device)
 
-crc_path = "/data/linears/logs/dino_vits16/offline/crc/20251110-164127567401_f6ae5d68/run_4/checkpoints/best.ckpt"
+crc_path = "./crc_best.ckpt"
 crc_linear = setup_linear_crc(crc_path).to(device)
 
 
@@ -296,4 +283,4 @@ demo = gr.TabbedInterface([breakhis, gleason, crc, bach],["BreakHis", "Gleason",
 
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()