Update app.py

app.py

@@ -2,22 +2,24 @@ import os
 import zipfile
 import numpy as np
 import torch
-import requests
 from transformers import ViTForImageClassification, AdamW
 import nibabel as nib
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
-from io import BytesIO
+import requests
+import tempfile
 
-# Function to download and extract zip file from URL
-def extract_zip_from_url(url, extract_to):
+# Function to download zip files from URL
+def download_zip(url, download_path):
     response = requests.get(url)
-    if response.status_code == 200:
-        with zipfile.ZipFile(BytesIO(response.content)) as zip_ref:
-            zip_ref.extractall(extract_to)
-    else:
-        raise ValueError(f"Unable to download zip file: {url}")
+    with open(download_path, 'wb') as file:
+        file.write(response.content)
+
+# Function to extract zip files
+def extract_zip(zip_file, extract_to):
+    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+        zip_ref.extractall(extract_to)
 
 # Preprocess images
 def preprocess_image(image_path):
@@ -45,13 +47,21 @@ def preprocess_image(image_path):
 def prepare_dataset(extracted_folder):
     image_paths = []
     labels = []
-    for disease_folder in ['alzheimers_dataset', 'parkinsons_dataset', 'MSjpg']:
-        folder_path = os.path.join(extracted_folder, disease_folder)
-        label = {'alzheimers_dataset': 0, 'parkinsons_dataset': 1, 'MSjpg': 2}[disease_folder]
+    
+    # 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 img_file in os.listdir(folder_path):
             if img_file.endswith(('.nii', '.jpg', '.jpeg')):
                 image_paths.append(os.path.join(folder_path, img_file))
                 labels.append(label)
+    
     return image_paths, labels
 
 # Custom Dataset class
@@ -92,23 +102,29 @@ def fine_tune_model(train_loader):
 # Streamlit UI for Fine-tuning
 st.title("Fine-tune ViT on MRI/CT Scans for MS & Neurodegenerative Diseases")
 
-# Input zip file URL
-zip_file_url = st.text_input("https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/neuroniiimages.zip")
+zip_url = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/neuroniiimages.zip"
+
 if st.button("Start Training"):
-    extraction_dir = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/extracttedfiles"
+    extraction_dir = "extracted_files"
     os.makedirs(extraction_dir, exist_ok=True)
 
-    # Extract the zip file from URL
-    extract_zip_from_url(zip_file_url, extraction_dir)
+    # 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)
+
+        # Extract the zip file
+        extract_zip(tmp_file.name, extraction_dir)
 
     # Prepare dataset
     image_paths, labels = prepare_dataset(extraction_dir)
     dataset = CustomImageDataset(image_paths, labels)
-    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}")
-
+    
+    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}")