Update app.py

app.py

@@ -7,39 +7,27 @@ import nibabel as nib
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
-import requests
 
-# 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:
+# 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)
 
-# 2. Preprocess images
+# Preprocess images
 def preprocess_image(image_path):
     ext = os.path.splitext(image_path)[-1].lower()
 
-    if ext == '.nii' or ext == '.nii.gz':
+    if ext in ['.nii', '.nii.gz']:
         nii_image = nib.load(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)  # Add channel dimension for MRI
+            image_tensor = image_tensor.unsqueeze(0)
 
     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()  # Change to [C, H, W]
+        image_tensor = torch.tensor(img_np).permute(2, 0, 1).float()
 
     else:
         raise ValueError(f"Unsupported format: {ext}")
@@ -47,21 +35,20 @@ def preprocess_image(image_path):
     image_tensor /= 255.0  # Normalize to [0, 1]
     return image_tensor
 
-# 3. Label images
+# Prepare dataset
 def prepare_dataset(extracted_folder):
     image_paths = []
     labels = []
     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]
-        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)
+        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
+# Custom Dataset class
 class CustomImageDataset(Dataset):
     def __init__(self, image_paths, labels):
         self.image_paths = image_paths
@@ -75,7 +62,7 @@ class CustomImageDataset(Dataset):
         label = self.labels[idx]
         return image, label
 
-# 5. Training function
+# Training function
 def fine_tune_model(train_loader):
     model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=3)
     model.train()
@@ -96,32 +83,28 @@ def fine_tune_model(train_loader):
             running_loss += loss.item()
     return running_loss / len(train_loader)
 
-# Streamlit UI
-st.title("Fine-tune ViT on MRI Scans")
+# Streamlit UI for Fine-tuning
+st.title("Fine-tune ViT on MRI/CT Scans for MS & Neurodegenerative Diseases")
 
-# 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")
+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"
 
 if st.button("Start Training"):
-    # Define an extraction directory
-    extraction_dir = 'extracted_files'
+    extraction_dir = "extracted_files"
     os.makedirs(extraction_dir, exist_ok=True)
 
-    # 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.")
+    # Extract both zip files
+    extract_zip(zip_file_1, extraction_dir)
+    extract_zip(zip_file_2, extraction_dir)
 
     # 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}")
 
+
+