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}")