Update app.py

app.py

@@ -2,8 +2,7 @@ import os
 import zipfile
 import numpy as np
 import torch
-from transformers import ViTForImageClassification
-from torch.optim import AdamW
+from transformers import SegformerForImageSegmentation, ResNetForImageClassification, AdamW
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
@@ -17,14 +16,12 @@ def extract_zip(zip_file, extract_to):
 def preprocess_image(image_path):
     ext = os.path.splitext(image_path)[-1].lower()
 
-    if ext in ['.npy']:
+    if ext == '.npy':
         image_data = np.load(image_path)
         image_tensor = torch.tensor(image_data).float()
-        if len(image_tensor.shape) == 2:  # If the image is 2D (grayscale)
-            image_tensor = image_tensor.unsqueeze(0)  # Add channel dimension
-        elif len(image_tensor.shape) == 3:  # If the image is 3D (height, width, channels)
-            image_tensor = image_tensor.permute(2, 0, 1).float()  # Change to (C, H, W)
-    
+        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)
@@ -38,7 +35,6 @@ 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):
@@ -50,11 +46,9 @@ def prepare_dataset(extracted_folder):
     for disease_folder in ['alzheimers_dataset', 'parkinsons_dataset', 'MSjpg']:
         folder_path = os.path.join(neuronii_path, disease_folder)
         
-        # Check if the subfolder exists
         if not os.path.exists(folder_path):
             print(f"Folder not found: {folder_path}")
-            continue  # Skip this folder if it's not found
-            
+            continue  
         label = {'alzheimers_dataset': 0, 'parkinsons_dataset': 1, 'MSjpg': 2}[disease_folder]
         
         for img_file in os.listdir(folder_path):
@@ -63,7 +57,6 @@ def prepare_dataset(extracted_folder):
                 labels.append(label)
             else:
                 print(f"Unsupported file: {img_file}")
-    
     print(f"Total images loaded: {len(image_paths)}")
     return image_paths, labels
 
@@ -81,12 +74,11 @@ class CustomImageDataset(Dataset):
         label = self.labels[idx]
         return image, label
 
-# Training function
-# Training function
-def fine_tune_model(train_loader):
-    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=3)
+# Training function for classification
+def fine_tune_classification_model(train_loader):
+    model = ResNetForImageClassification.from_pretrained('microsoft/resnet-50', num_labels=3)
     model.train()
-    optimizer = AdamW(model.parameters(), lr=1e-4)  # Use PyTorch's AdamW
+    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)
@@ -104,16 +96,15 @@ def fine_tune_model(train_loader):
     return running_loss / len(train_loader)
 
 # Streamlit UI for Fine-tuning
-st.title("Fine-tune ViT on MRI/CT Scans for MS & Neurodegenerative Diseases")
+st.title("Fine-tune ResNet for MRI/CT Scans Classification")
 
-# 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 (this is a placeholder; use requests or any other method to download the zip file)
+    # Download the zip file (placeholder)
     zip_file = "neuroniiimages.zip"  # Assuming you downloaded it with this name
 
     # Extract zip file
@@ -124,6 +115,37 @@ if st.button("Start Training"):
     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)
+    # Fine-tune the classification model
+    final_loss = fine_tune_classification_model(train_loader)
     st.write(f"Training Complete with Final Loss: {final_loss}")
+
+# Segmentation function (using SegFormer)
+def fine_tune_segmentation_model(train_loader):
+    model = SegformerForImageSegmentation.from_pretrained('nvidia/segformer-b0', 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)
+
+# Add a button for segmentation training
+if st.button("Start Segmentation Training"):
+    # Assuming the dataset for segmentation is prepared similarly
+    seg_train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
+    
+    # Fine-tune the segmentation model
+    final_loss_seg = fine_tune_segmentation_model(seg_train_loader)
+    st.write(f"Segmentation Training Complete with Final Loss: {final_loss_seg}")
+