Update app.py

app.py

@@ -8,7 +8,7 @@ from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
 
-# Function to extract zip files
+# Function to extract zip file
 def extract_zip(zip_file, extract_to):
     with zipfile.ZipFile(zip_file, 'r') as zip_ref:
         zip_ref.extractall(extract_to)
@@ -41,7 +41,7 @@ def prepare_dataset(extracted_folder):
     labels = []
     for disease_folder in ['alzheimers_dataset', 'parkinsons_dataset', 'ms']:
         folder_path = os.path.join(extracted_folder, disease_folder)
-        label = {'alzheimers': 0, 'parkinsons': 1, 'ms': 2}[disease_folder]
+        label = {'alzheimers_dataset': 0, 'parkinsons_dataset': 1, 'ms': 2}[disease_folder]
         for img_file in os.listdir(folder_path):
             if img_file.endswith(('.nii', '.jpg', '.jpeg')):
                 image_paths.append(os.path.join(folder_path, img_file))
@@ -86,16 +86,14 @@ def fine_tune_model(train_loader):
 # Streamlit UI for Fine-tuning
 st.title("Fine-tune ViT on MRI/CT Scans for MS & Neurodegenerative Diseases")
 
-zip_file_1 = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/archive%20(5).zip"
-zip_file_2 = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/MS.zip"
+zip_file = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/archive%20(5).zip"
 
 if st.button("Start Training"):
     extraction_dir = "extracted_files"
     os.makedirs(extraction_dir, exist_ok=True)
 
-    # Extract both zip files
-    extract_zip(zip_file_1, extraction_dir)
-    extract_zip(zip_file_2, extraction_dir)
+    # Extract the zip file
+    extract_zip(zip_file, extraction_dir)
 
     # Prepare dataset
     image_paths, labels = prepare_dataset(extraction_dir)