cd updates

README.md

@@ -13,14 +13,16 @@ app_port: 8501
 </p>
 
 <p align="center">
-  <a href="https://github.com/anirudhbalaraman/WSAttention-Prostate/actions"><img alt="Continuous Integration" src="https://github.com/anirudhbalaraman/WSAttention-Prostate/actions/workflows/ci.yaml/badge.svg"></a>
+  <a href="https://huggingface.co/spaces/anirudh0410/WSA_Prostate">
+    <img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue" alt="Hugging Face Spaces">
+  </a>
+<a href="https://github.com/anirudhbalaraman/WSAttention-Prostate/actions/workflows/ci.yml">
+    <img src="https://github.com/anirudhbalaraman/WSAttention-Prostate/actions/workflows/ci.yml/badge.svg" alt="CI/CD Status">
+  </a>
   <img src="https://img.shields.io/badge/python-3.9-blue?logo=python&logoColor=white" alt="Python 3.9">
-  <img src="https://img.shields.io/badge/pytorch-2.5-ee4c2c?logo=pytorch&logoColor=white" alt="PyTorch 2.5">
-  <img src="https://img.shields.io/badge/MONAI-1.4-3ddc84" alt="MONAI 1.4">
+  <img src="https://img.shields.io/badge/pytorch-2.5-ee4c2c?logo=pytorch&logoColor=white" alt="PyTorch 2.5">  
+  <img src="https://img.shields.io/badge/docker-automated-blue?logo=docker&logoColor=white" alt="Docker">
   <img src="https://img.shields.io/badge/license-Apache%202.0-green" alt="License">
-  <a href="https://github.com/astral-sh/ruff">
-  <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Code style: ruff">
-  </a>
 </p>
 
 # Weakly Supervised Attention-Based Deep Learning for Prostate Cancer Characterization from Bi-Parametric Prostate MRI.

run_inference.py

@@ -31,9 +31,9 @@ def load_pirads_model(num_classes, mil_mode, project_dir, device):
     model.eval()  # Set to evaluation mode
     return model
 @st.cache_resource
-def load_cspca_model(pirads_model, project_dir, device):
+def load_cspca_model(_pirads_model, project_dir, device):
     # Move the model initialization inside here
-    model = CSPCAModel(backbone=pirads_model).to(device)
+    model = CSPCAModel(backbone=_pirads_model).to(device)
     checkpt = torch.load(
         os.path.join(project_dir, "models", "cspca_model.pth"), map_location="cpu"
     )