app.py

import os
import zipfile
import numpy as np
import torch
from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor
from transformers import ResNetForImageClassification, AdamW
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import streamlit as st

# Load feature extractor and model
feature_extractor = SegformerFeatureExtractor.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512')
segformer_model = SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512')

# 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 = []
    
    # Define the mapping of folders to labels
    folder_label_mapping = {
        'alzheimers_dataset': 0,
        'parkinsons_dataset': 1,
        'MSjpg': 2
    }

    for disease_folder, label in folder_label_mapping.items():
        folder_path = os.path.join(neuronii_path, disease_folder)
        
        if not os.path.exists(folder_path):
            print(f"Folder not found: {folder_path}")
            continue  
        
        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):
    # Load the ResNet model with ignore_mismatched_sizes
    model = ResNetForImageClassification.from_pretrained('microsoft/resnet-50', num_labels=3, ignore_mismatched_sizes=True)
    
    # Update the classifier layer to match the number of labels
    if hasattr(model, 'classifier'):
        if isinstance(model.classifier, torch.nn.Sequential):
            model.classifier[-1] = torch.nn.Linear(model.classifier[-1].in_features, 3)  # Assuming 3 output classes
        else:
            model.classifier = torch.nn.Linear(model.classifier.in_features, 3)  # In case it's a Linear layer directly
    else:
        print("Classifier layer not found")

    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):  # Adjust epochs as needed
        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/Segmentation_and_classification/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):
    # Load the Segformer model with ignore_mismatched_sizes
    model = SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0', num_labels=3, ignore_mismatched_sizes=True)
    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):  # Adjust epochs as needed
        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}")