Update demo.py

demo.py

@@ -26,12 +26,15 @@ image_files = None
 selectedID = 0
 question_dropdown = None
 
+#NO GPU is available
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
 def seed_everything(seed=27):
     torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
+    #torch.cuda.manual_seed_all(seed)
     os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
+    #torch.backends.cudnn.deterministic = True
+    #torch.backends.cudnn.benchmark = False
 
 def load_visualbert_model(tokenizer, device, num_class=51, encoder_layers=6, n_heads=8, dropout=0.1, emb_dim=300):
     """
@@ -43,7 +46,7 @@ def load_visualbert_model(tokenizer, device, num_class=51, encoder_layers=6, n_h
         n_heads=n_heads,
         num_class=num_class,
     )
-    checkpoint = torch.load("checkpoint.tar", map_location=device)
+    checkpoint = torch.load("./checkpoint.tar", map_location=device)
     model.load_state_dict(checkpoint["model"])
     model.to(device)
     model.eval()
@@ -55,7 +58,7 @@ def load_surgvlp_encoder(device):
     """
     config_path = './utils/config_surgvlp.py'
     configs = Config.fromfile(config_path)['config']
-    encoder_model, encoder_preprocess = surgvlp.load(configs.model_config, device=device, pretrain='SurgVLP2.pth')
+    encoder_model, encoder_preprocess = surgvlp.load(configs.model_config, device=device, pretrain='./SurgVLP2.pth')
     encoder_model.eval()
     return encoder_model, encoder_preprocess
 
@@ -73,17 +76,22 @@ LABEL_LIST = [
 
 def main():
     seed_everything()
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cpu"
     tokenizer = BertTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
     visualbert_model = load_visualbert_model(tokenizer, device)
     encoder_model, encoder_preprocess = load_surgvlp_encoder(device)
+
+    print("Models loaded successfully.")
     
     # Define the directories containing images and corresponding label files.
     global image_files
-    images_dir = "./test_data/images/VID01/"
+    images_dir = "./test_data/images/VID01"
     labels_dir = "./test_data/labels/VID01/"
     image_files = [os.path.join(images_dir, f) for f in sorted(os.listdir(images_dir)) if f.lower().endswith('.png')]
     random.shuffle(image_files)
+
+    print(f"Found {len(image_files)} images.")
+
     # Get first 20 images.
     image_files = image_files[:20]
 
@@ -196,7 +204,9 @@ def main():
             inputs=[image_gallery, question_dropdown], 
             outputs=predictions_output
         )
+
+    print("Launching the Gradio UI...")
     demo.launch()
 
 if __name__ == "__main__":
-    main() 
+    main()