Update app.py

app.py

@@ -1,119 +1,23 @@
 import os
 import zipfile
-import requests
 import numpy as np
 import torch
-from transformers import ViTForImageClassification, AdamW
-import nibabel as nib
-from PIL import Image
-from torch.utils.data import Dataset, DataLoader
-import streamlit as st
-
-# Function to download the zip file from the URL
-def download_zip(url, save_path):
-    response = requests.get(url)
-    with open(save_path, 'wb') as f:
-        f.write(response.content)
-
-# 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)
-
-# Preprocess images
-def preprocess_image(image_path):
-    ext = os.path.splitext(image_path)[-1].lower()
-
-    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)
-
-    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()
-
-    else:
-        raise ValueError(f"Unsupported format: {ext}")
-
-    image_tensor /= 255.0  # Normalize to [0, 1]
-    return image_tensor
-
-# Prepare dataset
-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]
-        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
-class CustomImageDataset(Dataset):
-    def __init__(self, image_paths, labels):
-        self.image_paths = image_paths
-        self.labels = labels
-
-    def __len__(self):
-        return len(self.image_paths)
-
-    def __getitem__(self, idx):
-        image = preprocess_image(self.image_paths[idx])
-        label = self.labels[idx]
-        return image, label
-
-# Training function
-def fine_tune_model(train_loader):
-    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=3)
-    model.train()
-    optimizer = AdamW(model.parameters(), lr=1e-4)
-    criterion = torch.nn.CrossEntropyLoss()
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model.to(device)
-    
-    for epoch in range(10):
-        running_loss = 0.0
-        for images, labels in train_loader:
-            images, labels = images.to(device), labels.to(device)
-            optimizer.zero_grad()
-            outputs = model(pixel_values=images).logits
-            loss = criterion(outputs, labels)
-            loss.backward()
-            optimizer.step()
-            running_loss += loss.item()
-    return running_loss / len(train_loader)
-
-# Streamlit UI for Fine-tuning
-st.title("Fine-tune ViT on MRI/CT Scans for MS & Neurodegenerative Diseases")
-
-zip_file_url = import os
-import zipfile
 import requests
-import numpy as np
-import torch
 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
 
-# Function to download the zip file from the URL
-def download_zip(url, save_path):
+# Function to download and extract zip file from URL
+def extract_zip_from_url(url, extract_to):
     response = requests.get(url)
-    with open(save_path, 'wb') as f:
-        f.write(response.content)
-
-# 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)
+    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}")
 
 # Preprocess images
 def preprocess_image(image_path):
@@ -188,20 +92,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_url = "https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/neuroniiimages.zip"
-
+# Input zip file URL
+zip_file_url = st.text_input("https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/neuroniiimages.zip")
 if st.button("Start Training"):
     extraction_dir = "extracted_files"
-    zip_file_path = "archive_5.zip"
     os.makedirs(extraction_dir, exist_ok=True)
 
-    # Download the zip file
-    st.write("Downloading the zip file...")
-    download_zip(zip_file_url, zip_file_path)
-
-    # Extract the zip file
-    st.write("Extracting files...")
-    extract_zip(zip_file_path, extraction_dir)
+    # Extract the zip file from URL
+    extract_zip_from_url(zip_file_url, extraction_dir)
 
     # Prepare dataset
     image_paths, labels = prepare_dataset(extraction_dir)
@@ -209,31 +107,8 @@ if st.button("Start Training"):
     train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
 
     # Fine-tune the model
-    st.write("Fine-tuning the model...")
     final_loss = fine_tune_model(train_loader)
     st.write(f"Training Complete with Final Loss: {final_loss}")
 
-if st.button("Start Training"):
-    extraction_dir = "extracted_files"
-    zip_file_path = "archive_5.zip"
-    os.makedirs(extraction_dir, exist_ok=True)
-
-    # Download the zip file
-    st.write("Downloading the zip file...")
-    download_zip(zip_file_url, zip_file_path)
-
-    # Extract the zip file
-    st.write("Extracting files...")
-    extract_zip(zip_file_path, 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
-    st.write("Fine-tuning the model...")
-    final_loss = fine_tune_model(train_loader)
-    st.write(f"Training Complete with Final Loss: {final_loss}")