File size: 5,677 Bytes
edd91e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
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}")
|