Update app.py

app.py

@@ -7,14 +7,6 @@ import nibabel as nib
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
-import requests
-import tempfile
-
-# Function to download zip files from URL
-def download_zip(url, download_path):
-    response = requests.get(url)
-    with open(download_path, 'wb') as file:
-        file.write(response.content)
 
 # Function to extract zip files
 def extract_zip(zip_file, extract_to):
@@ -45,18 +37,24 @@ def preprocess_image(image_path):
 
 # Prepare dataset
 def prepare_dataset(extracted_folder):
+    # Ensure the path exists
+    neuronii_path = os.path.join(extracted_folder, "neuroniiimages")
+    
+    if not os.path.exists(neuronii_path):
+        raise FileNotFoundError(f"The folder neuroniiimages does not exist in the extracted folder: {neuronii_path}")
+    
     image_paths = []
     labels = []
     
-    # Define the paths for each disease dataset
-    datasets = {
-        'alzheimer_datasets': 0,
-        'parkinson_datasets': 1,
-        'MSjpg': 2
-    }
-    
-    for disease_folder, label in datasets.items():
-        folder_path = os.path.join(extracted_folder, 'neuroniiimages', disease_folder)
+    for disease_folder in ['alzheimers_datasets', 'parkinson_datasets', 'MSjpg']:
+        folder_path = os.path.join(neuronii_path, disease_folder)
+        
+        # Check if the subfolder exists
+        if not os.path.exists(folder_path):
+            raise FileNotFoundError(f"The folder {disease_folder} does not exist at path: {folder_path}")
+        
+        label = {'alzheimers_datasets': 0, 'parkinson_datasets': 1, 'MSjpg': 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))
@@ -102,29 +100,24 @@ 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_url = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/neuroniiimages.zip"
+# Provide the correct zip file URL
+zip_file_url = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/neuroniiimages.zip"
 
 if st.button("Start Training"):
     extraction_dir = "extracted_files"
     os.makedirs(extraction_dir, exist_ok=True)
 
-    # Download the zip file to a temporary file
-    with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as tmp_file:
-        download_zip(zip_url, tmp_file.name)
+    # Download the zip file (this is a placeholder; use requests or any other method to download the zip file)
+    zip_file = "neuroniiimages.zip"  # Assuming you downloaded it with this name
 
-        # Extract the zip file
-        extract_zip(tmp_file.name, extraction_dir)
+    # Extract zip file
+    extract_zip(zip_file, extraction_dir)
 
     # Prepare dataset
     image_paths, labels = prepare_dataset(extraction_dir)
     dataset = CustomImageDataset(image_paths, labels)
-    
-    if len(image_paths) == 0:
-        st.error("No images found in the specified directory. Please check the folder structure.")
-    else:
-        train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
-
-        # Fine-tune the model
-        final_loss = fine_tune_model(train_loader)
-        st.write(f"Training Complete with Final Loss: {final_loss}")
+    train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
 
+    # Fine-tune the model
+    final_loss = fine_tune_model(train_loader)
+    st.write(f"Training Complete with Final Loss: {final_loss}")