Update app.py

app.py

@@ -7,10 +7,22 @@ import nibabel as nib
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
+import requests
 
-# 1. Function to extract zip files
-def extract_zip(zip_file, extract_to):
-    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+# 1. Function to download and extract zip files
+def download_and_extract_zip(zip_url, extract_to):
+    local_zip_path = os.path.join(extract_to, os.path.basename(zip_url))
+    os.makedirs(extract_to, exist_ok=True)
+    
+    # Download the zip file
+    with requests.get(zip_url, stream=True) as r:
+        r.raise_for_status()
+        with open(local_zip_path, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+
+    # Extract the zip file
+    with zipfile.ZipFile(local_zip_path, 'r') as zip_ref:
         zip_ref.extractall(extract_to)
 
 # 2. Preprocess images
@@ -22,12 +34,12 @@ def preprocess_image(image_path):
         image_data = nii_image.get_fdata()
         image_tensor = torch.tensor(image_data).float()
         if len(image_tensor.shape) == 3:
-            image_tensor = image_tensor.unsqueeze(0)
+            image_tensor = image_tensor.unsqueeze(0)  # Add channel dimension for MRI
 
     elif ext in ['.jpg', '.jpeg']:
         img = Image.open(image_path).convert('RGB').resize((224, 224))
         img_np = np.array(img)
-        image_tensor = torch.tensor(img_np).permute(2, 0, 1).float()
+        image_tensor = torch.tensor(img_np).permute(2, 0, 1).float()  # Change to [C, H, W]
 
     else:
         raise ValueError(f"Unsupported format: {ext}")
@@ -42,10 +54,11 @@ def prepare_dataset(extracted_folder):
     for disease_folder in ['alzheimers', 'parkinsons', 'ms']:
         folder_path = os.path.join(extracted_folder, disease_folder)
         label = {'alzheimers': 0, 'parkinsons': 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))
-                labels.append(label)
+        if os.path.exists(folder_path):
+            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
 
 # 4. Custom Dataset
@@ -86,25 +99,29 @@ def fine_tune_model(train_loader):
 # Streamlit UI
 st.title("Fine-tune ViT on MRI Scans")
 
-# Input for zip file paths
-zip_file_1 = st.text_input("https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/archive%20(5).zip")
-zip_file_2 = st.text_input("https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/MS.zip")
+# Input for zip file URLs
+zip_file_1 = st.text_input("Enter URL for the 1st zip file:", "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/archive%20(5).zip")
+zip_file_2 = st.text_input("Enter URL for the 2nd zip file:", "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/MS.zip")
 
 if st.button("Start Training"):
     # Define an extraction directory
-    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 both zip files
-    extract_zip(zip_file_1, extraction_dir)
-    extract_zip(zip_file_2, extraction_dir)
+    # Download and extract both zip files
+    st.write("Downloading and extracting files...")
+    download_and_extract_zip(zip_file_1, extraction_dir)
+    download_and_extract_zip(zip_file_2, extraction_dir)
+    st.write("Extraction complete.")
 
     # Prepare dataset
+    st.write("Preparing 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
+    st.write("Training model...")
     final_loss = fine_tune_model(train_loader)
     st.write(f"Training Complete with Final Loss: {final_loss}")