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 import gradio as gr 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 import gradio as gr # 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 = [] 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): # Load the ResNet model with ignore_mismatched_sizes model = ResNetForImageClassification.from_pretrained('microsoft/resnet-50', num_labels=3, ignore_mismatched_sizes=True) # Print model architecture to identify the classifier layer print(model) # Inspect the model structure # Update the classifier layer to match the number of labels if hasattr(model, 'classifier'): model.classifier = torch.nn.Linear(model.classifier.in_features, 3) # Assuming 3 output classes else: # Access the linear layer differently if 'classifier' does not exist 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/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): 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}")