Create app.py

app.py

@@ -0,0 +1,154 @@
+import os
+import zipfile
+import numpy as np
+import torch
+from transformers import SegformerForImageSegmentation, ResNetForImageClassification, AdamW
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+import streamlit as st
+import gradio as gr
+
+# Load the Segformer model using Gradio (Optional)
+gr.load("models/nvidia/segformer-b0-finetuned-ade-512-512").launch()
+
+# 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):
+    ext = os.path.splitext(image_path)[-1].lower()
+
+    if ext == '.npy':
+        image_data = np.load(image_path)
+        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):
+    neuronii_path = os.path.join(extracted_folder, "neuroniiimages")
+    
+    if not os.path.exists(neuronii_path):
+        raise FileNotFoundError(f"The folder neuroniiimages does not exist in the extracted folder: {neuronii_path}")
+    
+    image_paths = []
+    labels = []
+    
+    for disease_folder in ['alzheimers_dataset', 'parkinsons_dataset', 'MSjpg']:
+        folder_path = os.path.join(neuronii_path, disease_folder)
+        
+        if not os.path.exists(folder_path):
+            print(f"Folder not found: {folder_path}")
+            continue  
+        label = {'alzheimers_dataset': 0, 'parkinsons_dataset': 1, 'MSjpg': 2}[disease_folder]
+        
+        for img_file in os.listdir(folder_path):
+            if img_file.endswith(('.npy', '.jpg', '.jpeg')):
+                image_paths.append(os.path.join(folder_path, img_file))
+                labels.append(label)
+            else:
+                print(f"Unsupported file: {img_file}")
+    print(f"Total images loaded: {len(image_paths)}")
+    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 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)
+    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 ResNet for MRI/CT Scans Classification")
+
+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 (placeholder)
+    zip_file = "neuroniiimages.zip"  # Assuming you downloaded it with this name
+
+    # Extract zip file
+    extract_zip(zip_file, 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 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}")