import os, torch, time, gc, datetime, cv2
torch.backends.cudnn.benchmark = True
import numpy as np
import pandas as pd
from torch.optim import Adagrad, AdamW
from torch.autograd import grad
from torch.optim.lr_scheduler import StepLR
import torch.nn.functional as F
import matplotlib.pyplot as plt
from PIL import Image
from sklearn.metrics import auc, precision_recall_curve
from IPython.display import display
from multiprocessing import cpu_count
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader
[docs]
def apply_model(src, model_path, image_size=224, batch_size=64, normalize=True, n_jobs=10):
from .io import NoClassDataset
from .utils import print_progress
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if normalize:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.CenterCrop(size=(image_size, image_size)),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])
else:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.CenterCrop(size=(image_size, image_size))])
model = torch.load(model_path)
print(model)
print(f'Loading dataset in {src} with {len(src)} images')
dataset = NoClassDataset(data_dir=src, transform=transform, shuffle=True, load_to_memory=False)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=n_jobs)
print(f'Loaded {len(src)} images')
result_loc = os.path.splitext(model_path)[0]+datetime.date.today().strftime('%y%m%d')+'_'+os.path.splitext(model_path)[1]+'_test_result.csv'
print(f'Results wil be saved in: {result_loc}')
model.eval()
model = model.to(device)
prediction_pos_probs = []
filenames_list = []
time_ls = []
with torch.no_grad():
for batch_idx, (batch_images, filenames) in enumerate(data_loader, start=1):
start = time.time()
images = batch_images.to(torch.float).to(device)
outputs = model(images)
batch_prediction_pos_prob = torch.sigmoid(outputs).cpu().numpy()
prediction_pos_probs.extend(batch_prediction_pos_prob.tolist())
filenames_list.extend(filenames)
stop = time.time()
duration = stop - start
time_ls.append(duration)
files_processed = batch_idx*batch_size
files_to_process = len(data_loader)
print_progress(files_processed, files_to_process, n_jobs=n_jobs, time_ls=time_ls, batch_size=batch_size, operation_type="Generating predictions")
data = {'path':filenames_list, 'pred':prediction_pos_probs}
df = pd.DataFrame(data, index=None)
df.to_csv(result_loc, index=True, header=True, mode='w')
torch.cuda.empty_cache()
torch.cuda.memory.empty_cache()
return df
[docs]
def apply_model_to_tar(settings={}):
from .io import TarImageDataset
from .utils import process_vision_results, print_progress
#if os.path.exists(settings['dataset']):
# tar_path = settings['dataset']
#else:
# tar_path = os.path.join(settings['src'], 'datasets', settings['dataset'])
tar_path = settings['tar_path']
model_path = settings['model_path']
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if settings['normalize']:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.CenterCrop(size=(settings['image_size'], settings['image_size'])),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])
else:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.CenterCrop(size=(settings['image_size'], settings['image_size']))])
if settings['verbose']:
print(f"Loading model from {model_path}")
print(f"Loading dataset from {tar_path}")
model = torch.load(settings['model_path'])
dataset = TarImageDataset(tar_path, transform=transform)
data_loader = DataLoader(dataset, batch_size=settings['batch_size'], shuffle=True, num_workers=settings['n_jobs'], pin_memory=True)
model_name = os.path.splitext(os.path.basename(model_path))[0]
#dataset_name = os.path.splitext(os.path.basename(settings['dataset']))[0]
dataset_name = os.path.splitext(os.path.basename(settings['tar_path']))[0]
date_name = datetime.date.today().strftime('%y%m%d')
dst = os.path.dirname(tar_path)
result_loc = f'{dst}/{date_name}_{dataset_name}_{model_name}_result.csv'
model.eval()
model = model.to(device)
if settings['verbose']:
print(model)
print(f'Generated dataset with {len(dataset)} images')
print(f'Generating loader from {len(data_loader)} batches')
print(f'Results wil be saved in: {result_loc}')
print(f'Model is in eval mode')
print(f'Model loaded to device')
prediction_pos_probs = []
filenames_list = []
time_ls = []
gc.collect()
with torch.no_grad():
for batch_idx, (batch_images, filenames) in enumerate(data_loader, start=1):
start = time.time()
images = batch_images.to(torch.float).to(device)
outputs = model(images)
batch_prediction_pos_prob = torch.sigmoid(outputs).cpu().numpy()
prediction_pos_probs.extend(batch_prediction_pos_prob.tolist())
filenames_list.extend(filenames)
stop = time.time()
duration = stop - start
time_ls.append(duration)
files_processed = batch_idx*settings['batch_size']
files_to_process = len(data_loader)*settings['batch_size']
print_progress(files_processed, files_to_process, n_jobs=settings['n_jobs'], time_ls=time_ls, batch_size=settings['batch_size'], operation_type="Tar dataset")
data = {'path':filenames_list, 'pred':prediction_pos_probs}
df = pd.DataFrame(data, index=None)
df = process_vision_results(df, settings['score_threshold'])
df.to_csv(result_loc, index=True, header=True, mode='w')
print(f"Saved results to {result_loc}")
torch.cuda.empty_cache()
torch.cuda.memory.empty_cache()
return df
[docs]
def test_model_core(model, loader, loader_name, epoch, loss_type):
from .utils import calculate_loss, classification_metrics
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.eval()
loss = 0
correct = 0
total_samples = 0
prediction_pos_probs = []
all_labels = []
filenames = []
true_targets = []
predicted_outputs = []
model = model.to(device)
with torch.no_grad():
for batch_idx, (data, target, filename) in enumerate(loader, start=1): # Assuming loader provides filenames
start_time = time.time()
data, target = data.to(device), target.to(device).float()
output = model(data)
loss += F.binary_cross_entropy_with_logits(output, target, reduction='sum').item()
loss = calculate_loss(output, target, loss_type=loss_type)
loss += loss.item()
total_samples += data.size(0)
pred = torch.where(output >= 0.5,
torch.Tensor([1.0]).to(device).float(),
torch.Tensor([0.0]).to(device).float())
correct += pred.eq(target.view_as(pred)).sum().item()
batch_prediction_pos_prob = torch.sigmoid(output).cpu().numpy()
prediction_pos_probs.extend(batch_prediction_pos_prob.tolist())
all_labels.extend(target.cpu().numpy().tolist())
# Storing intermediate results in lists
true_targets.extend(target.cpu().numpy().tolist())
predicted_outputs.extend(pred.cpu().numpy().tolist())
filenames.extend(filename)
mean_loss = loss / total_samples
acc = correct / total_samples
end_time = time.time()
test_time = end_time - start_time
#print(f'\rTest: epoch: {epoch} Accuracy: {acc:.5f} batch: {batch_idx}/{len(loader)} loss: {mean_loss:.5f} time {test_time:.5f}', end='\r', flush=True)
# Constructing the DataFrame
results_df = pd.DataFrame({
'filename': filenames,
'true_label': true_targets,
'predicted_label': predicted_outputs,
'class_1_probability':prediction_pos_probs})
loss /= len(loader)
data_df = classification_metrics(all_labels, prediction_pos_probs, loss, epoch)
return data_df, prediction_pos_probs, all_labels, results_df
[docs]
def train_test_model(settings):
from .io import _copy_missclassified
from .utils import pick_best_model, save_settings
from .io import generate_loaders
from .settings import get_train_test_model_settings
settings = get_train_test_model_settings(settings)
torch.cuda.empty_cache()
torch.cuda.memory.empty_cache()
gc.collect()
src = settings['src']
channels_str = ''.join(settings['train_channels'])
dst = os.path.join(src,'model', settings['model_type'], channels_str, str(f"epochs_{settings['epochs']}"))
os.makedirs(dst, exist_ok=True)
settings['src'] = src
settings['dst'] = dst
if settings['custom_model']:
model = torch.load(settings['custom_model'])
if settings['train']:
if settings['train'] and settings['test']:
save_settings(settings, name=f"train_test_{settings['model_type']}_{settings['epochs']}", show=True)
elif settings['train'] is True:
save_settings(settings, name=f"train_{settings['model_type']}_{settings['epochs']}", show=True)
elif settings['test'] is True:
save_settings(settings, name=f"test_{settings['model_type']}_{settings['epochs']}", show=True)
if settings['train']:
train, val, train_fig = generate_loaders(src,
mode='train',
image_size=settings['image_size'],
batch_size=settings['batch_size'],
classes=settings['classes'],
n_jobs=settings['n_jobs'],
validation_split=settings['val_split'],
pin_memory=settings['pin_memory'],
normalize=settings['normalize'],
channels=settings['train_channels'],
augment=settings['augment'],
verbose=settings['verbose'])
#train_batch_1_figure = os.path.join(dst, 'batch_1.pdf')
#train_fig.savefig(train_batch_1_figure, format='pdf', dpi=300)
if settings['train']:
model, model_path = train_model(dst = settings['dst'],
model_type=settings['model_type'],
train_loaders = train,
epochs = settings['epochs'],
learning_rate = settings['learning_rate'],
init_weights = settings['init_weights'],
weight_decay = settings['weight_decay'],
amsgrad = settings['amsgrad'],
optimizer_type = settings['optimizer_type'],
use_checkpoint = settings['use_checkpoint'],
dropout_rate = settings['dropout_rate'],
n_jobs = settings['n_jobs'],
val_loaders = val,
test_loaders = None,
intermedeate_save = settings['intermedeate_save'],
schedule = settings['schedule'],
loss_type=settings['loss_type'],
gradient_accumulation=settings['gradient_accumulation'],
gradient_accumulation_steps=settings['gradient_accumulation_steps'],
channels=settings['train_channels'])
if settings['test']:
test, _, train_fig = generate_loaders(src,
mode='test',
image_size=settings['image_size'],
batch_size=settings['batch_size'],
classes=settings['classes'],
n_jobs=settings['n_jobs'],
validation_split=0.0,
pin_memory=settings['pin_memory'],
normalize=settings['normalize'],
channels=settings['train_channels'],
augment=False,
verbose=settings['verbose'])
if model == None:
model_path = pick_best_model(src+'/model')
print(f'Best model: {model_path}')
model = torch.load(model_path, map_location=lambda storage, loc: storage)
model_type = settings['model_type']
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(type(model))
print(model)
model_fldr = dst
time_now = datetime.date.today().strftime('%y%m%d')
result_loc = f"{model_fldr}/{settings['model_type']}_time_{time_now}_test_result.csv"
acc_loc = f"{model_fldr}/{settings['model_type']}_time_{time_now}_test_acc.csv"
print(f'Results wil be saved in: {result_loc}')
result, accuracy = test_model_performance(loaders=test,
model=model,
loader_name_list='test',
epoch=1,
loss_type=settings['loss_type'])
result.to_csv(result_loc, index=True, header=True, mode='w')
accuracy.to_csv(acc_loc, index=True, header=True, mode='w')
_copy_missclassified(accuracy)
torch.cuda.empty_cache()
torch.cuda.memory.empty_cache()
gc.collect()
if settings['train']:
return model_path
if settings['test']:
return result_loc
[docs]
def train_model(dst, model_type, train_loaders, epochs=100, learning_rate=0.0001, weight_decay=0.05, amsgrad=False, optimizer_type='adamw', use_checkpoint=False, dropout_rate=0, n_jobs=20, val_loaders=None, test_loaders=None, init_weights='imagenet', intermedeate_save=None, chan_dict=None, schedule = None, loss_type='binary_cross_entropy_with_logits', gradient_accumulation=False, gradient_accumulation_steps=4, channels=['r','g','b'], verbose=False):
"""
Trains a model using the specified parameters.
Args:
dst (str): The destination path to save the model and results.
model_type (str): The type of model to train.
train_loaders (list): A list of training data loaders.
epochs (int, optional): The number of training epochs. Defaults to 100.
learning_rate (float, optional): The learning rate for the optimizer. Defaults to 0.0001.
weight_decay (float, optional): The weight decay for the optimizer. Defaults to 0.05.
amsgrad (bool, optional): Whether to use AMSGrad for the optimizer. Defaults to False.
optimizer_type (str, optional): The type of optimizer to use. Defaults to 'adamw'.
use_checkpoint (bool, optional): Whether to use checkpointing during training. Defaults to False.
dropout_rate (float, optional): The dropout rate for the model. Defaults to 0.
n_jobs (int, optional): The number of n_jobs for data loading. Defaults to 20.
val_loaders (list, optional): A list of validation data loaders. Defaults to None.
test_loaders (list, optional): A list of test data loaders. Defaults to None.
init_weights (str, optional): The initialization weights for the model. Defaults to 'imagenet'.
intermedeate_save (list, optional): The intermediate save thresholds. Defaults to None.
chan_dict (dict, optional): The channel dictionary. Defaults to None.
schedule (str, optional): The learning rate schedule. Defaults to None.
loss_type (str, optional): The loss function type. Defaults to 'binary_cross_entropy_with_logits'.
gradient_accumulation (bool, optional): Whether to use gradient accumulation. Defaults to False.
gradient_accumulation_steps (int, optional): The number of steps for gradient accumulation. Defaults to 4.
Returns:
None
"""
from .io import _save_model, _save_progress
from .utils import calculate_loss, choose_model
print(f'Train batches:{len(train_loaders)}, Validation batches:{len(val_loaders)}')
if test_loaders != None:
print(f'Test batches:{len(test_loaders)}')
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print(f'Using {device} for Torch')
kwargs = {'n_jobs': n_jobs, 'pin_memory': True} if use_cuda else {}
model = choose_model(model_type, device, init_weights, dropout_rate, use_checkpoint, verbose=verbose)
if model is None:
print(f'Model {model_type} not found')
return
print(f'Loading Model to {device}...')
model.to(device)
if optimizer_type == 'adamw':
optimizer = AdamW(model.parameters(), lr=learning_rate, betas=(0.9, 0.999), weight_decay=weight_decay, amsgrad=amsgrad)
if optimizer_type == 'adagrad':
optimizer = Adagrad(model.parameters(), lr=learning_rate, eps=1e-8, weight_decay=weight_decay)
if schedule == 'step_lr':
StepLR_step_size = int(epochs/5)
StepLR_gamma = 0.75
scheduler = StepLR(optimizer, step_size=StepLR_step_size, gamma=StepLR_gamma)
elif schedule == 'reduce_lr_on_plateau':
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, verbose=True)
else:
scheduler = None
time_ls = []
# Initialize lists to accumulate results
accumulated_train_dicts = []
accumulated_val_dicts = []
accumulated_test_dicts = []
print(f'Training ...')
for epoch in range(1, epochs+1):
model.train()
start_time = time.time()
running_loss = 0.0
# Initialize gradients if using gradient accumulation
if gradient_accumulation:
optimizer.zero_grad()
for batch_idx, (data, target, filenames) in enumerate(train_loaders, start=1):
data, target = data.to(device), target.to(device).float()
output = model(data)
loss = calculate_loss(output, target, loss_type=loss_type)
# Normalize loss if using gradient accumulation
if gradient_accumulation:
loss /= gradient_accumulation_steps
running_loss += loss.item() * gradient_accumulation_steps # correct the running_loss
loss.backward()
# Step optimizer if not using gradient accumulation or every gradient_accumulation_steps
if not gradient_accumulation or (batch_idx % gradient_accumulation_steps == 0):
optimizer.step()
optimizer.zero_grad()
avg_loss = running_loss / batch_idx
batch_size = len(train_loaders)
duration = time.time() - start_time
time_ls.append(duration)
#print(f'Progress: {batch_idx}/{batch_size}, operation_type: DL-Batch, Epoch {epoch}/{epochs}, Loss {avg_loss}, Time {duration}')
end_time = time.time()
train_time = end_time - start_time
train_dict, _ = evaluate_model_performance(model, train_loaders, epoch, loss_type=loss_type)
train_dict['train_time'] = train_time
accumulated_train_dicts.append(train_dict)
if val_loaders != None:
val_dict, _ = evaluate_model_performance(model, val_loaders, epoch, loss_type=loss_type)
accumulated_val_dicts.append(val_dict)
if schedule == 'reduce_lr_on_plateau':
val_loss = val_dict['loss']
print(f"Progress: {train_dict['epoch']}/{epochs}, operation_type: Training, Train Loss: {train_dict['loss']:.3f}, Val Loss: {val_dict['loss']:.3f}, Train acc.: {train_dict['accuracy']:.3f}, Val acc.: {val_dict['accuracy']:.3f}, Train NC acc.: {train_dict['neg_accuracy']:.3f}, Val NC acc.: {val_dict['neg_accuracy']:.3f}, Train PC acc.: {train_dict['pos_accuracy']:.3f}, Val PC acc.: {val_dict['pos_accuracy']:.3f}, Train PRAUC: {train_dict['prauc']:.3f}, Val PRAUC: {val_dict['prauc']:.3f}")
else:
print(f"Progress: {train_dict['epoch']}/{epochs}, operation_type: Training, Train Loss: {train_dict['loss']:.3f}, Train acc.: {train_dict['accuracy']:.3f}, Train NC acc.: {train_dict['neg_accuracy']:.3f}, Train PC acc.: {train_dict['pos_accuracy']:.3f}, Train PRAUC: {train_dict['prauc']:.3f}")
if test_loaders != None:
test_dict, _ = evaluate_model_performance(model, test_loaders, epoch, loss_type=loss_type)
accumulated_test_dicts.append(test_dict)
print(f"Progress: {test_dict['epoch']}/{epochs}, operation_type: Training, Train Loss: {test_dict['loss']:.3f}, Train acc.: {test_dict['accuracy']:.3f}, Train NC acc.: {test_dict['neg_accuracy']:.3f}, Train PC acc.: {test_dict['pos_accuracy']:.3f}, Train PRAUC: {test_dict['prauc']:.3f}")
if scheduler:
if schedule == 'reduce_lr_on_plateau':
scheduler.step(val_loss)
if schedule == 'step_lr':
scheduler.step()
if accumulated_train_dicts and accumulated_val_dicts:
train_df = pd.DataFrame(accumulated_train_dicts)
validation_df = pd.DataFrame(accumulated_val_dicts)
_save_progress(dst, train_df, validation_df)
accumulated_train_dicts, accumulated_val_dicts = [], []
elif accumulated_train_dicts:
train_df = pd.DataFrame(accumulated_train_dicts)
_save_progress(dst, train_df, None)
accumulated_train_dicts = []
elif accumulated_test_dicts:
test_df = pd.DataFrame(accumulated_test_dicts)
_save_progress(dst, test_df, None)
accumulated_test_dicts = []
batch_size = len(train_loaders)
duration = time.time() - start_time
time_ls.append(duration)
model_path = _save_model(model, model_type, train_dict, dst, epoch, epochs, intermedeate_save=[0.99,0.98,0.95,0.94], channels=channels)
return model, model_path
[docs]
def generate_activation_map(settings):
from .utils import SaliencyMapGenerator, GradCAMGenerator, SelectChannels, activation_maps_to_database, activation_correlations_to_database
from .utils import print_progress, save_settings, calculate_activation_correlations
from .io import TarImageDataset
from .settings import get_default_generate_activation_map_settings
torch.cuda.empty_cache()
gc.collect()
plt.clf()
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
source_folder = os.path.dirname(os.path.dirname(settings['dataset']))
settings['src'] = source_folder
settings = get_default_generate_activation_map_settings(settings)
save_settings(settings, name=f"{settings['cam_type']}_settings", show=False)
if settings['model_type'] == 'maxvit' and settings['target_layer'] == None:
settings['target_layer'] = 'base_model.blocks.3.layers.1.layers.MBconv.layers.conv_b'
if settings['cam_type'] in ['saliency_image', 'saliency_channel']:
settings['target_layer'] = None
# Set number of jobs for loading
n_jobs = settings['n_jobs']
if n_jobs is None:
n_jobs = max(1, cpu_count() - 4)
# Set transforms for images
transform = transforms.Compose([
transforms.ToTensor(),
transforms.CenterCrop(size=(settings['image_size'], settings['image_size'])),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) if settings['normalize_input'] else None,
SelectChannels(settings['channels'])
])
# Handle dataset path
if not os.path.exists(settings['dataset']):
print(f"Dataset not found at {settings['dataset']}")
return
# Load the model
model = torch.load(settings['model_path'])
model.to(device)
model.eval()
# Create directory for saving activation maps if it does not exist
dataset_dir = os.path.dirname(settings['dataset'])
dataset_name = os.path.splitext(os.path.basename(settings['dataset']))[0]
save_dir = os.path.join(dataset_dir, dataset_name, settings['cam_type'])
batch_grid_fldr = os.path.join(save_dir, 'batch_grids')
if settings['save']:
os.makedirs(save_dir, exist_ok=True)
print(f"Activation maps will be saved in: {save_dir}")
if settings['plot']:
os.makedirs(batch_grid_fldr, exist_ok=True)
print(f"Batch grid maps will be saved in: {batch_grid_fldr}")
# Load dataset
dataset = TarImageDataset(settings['dataset'], transform=transform)
data_loader = DataLoader(dataset, batch_size=settings['batch_size'], shuffle=settings['shuffle'], num_workers=n_jobs, pin_memory=True)
# Initialize generator based on cam_type
if settings['cam_type'] in ['gradcam', 'gradcam_pp']:
cam_generator = GradCAMGenerator(model, target_layer=settings['target_layer'], cam_type=settings['cam_type'])
elif settings['cam_type'] in ['saliency_image', 'saliency_channel']:
cam_generator = SaliencyMapGenerator(model)
time_ls = []
for batch_idx, (inputs, filenames) in enumerate(data_loader):
start = time.time()
img_paths = []
inputs = inputs.to(device)
# Compute activation maps and predictions
if settings['cam_type'] in ['gradcam', 'gradcam_pp']:
activation_maps, predicted_classes = cam_generator.compute_gradcam_and_predictions(inputs)
elif settings['cam_type'] in ['saliency_image', 'saliency_channel']:
activation_maps, predicted_classes = cam_generator.compute_saliency_and_predictions(inputs)
# Move activation maps to CPU
activation_maps = activation_maps.cpu()
# Sum saliency maps for 'saliency_image' type
if settings['cam_type'] == 'saliency_image':
summed_activation_maps = []
for i in range(activation_maps.size(0)):
activation_map = activation_maps[i]
#print(f"1: {activation_map.shape}")
activation_map_sum = activation_map.sum(dim=0, keepdim=False)
#print(f"2: {activation_map.shape}")
activation_map_sum = np.squeeze(activation_map_sum, axis=0)
#print(f"3: {activation_map_sum.shape}")
summed_activation_maps.append(activation_map_sum)
activation_maps = torch.stack(summed_activation_maps)
# For plotting
if settings['plot']:
fig = cam_generator.plot_activation_grid(inputs, activation_maps, predicted_classes, overlay=settings['overlay'], normalize=settings['normalize'])
pdf_save_path = os.path.join(batch_grid_fldr,f"batch_{batch_idx}_grid.pdf")
fig.savefig(pdf_save_path, format='pdf')
print(f"Saved batch grid to {pdf_save_path}")
#plt.show()
display(fig)
for i in range(inputs.size(0)):
activation_map = activation_maps[i].detach().numpy()
if settings['cam_type'] in ['saliency_image', 'gradcam', 'gradcam_pp']:
#activation_map = activation_map.sum(axis=0)
activation_map = (activation_map - activation_map.min()) / (activation_map.max() - activation_map.min())
activation_map = (activation_map * 255).astype(np.uint8)
activation_image = Image.fromarray(activation_map, mode='L')
elif settings['cam_type'] == 'saliency_channel':
# Handle each channel separately and save as RGB
rgb_activation_map = np.zeros((activation_map.shape[1], activation_map.shape[2], 3), dtype=np.uint8)
for c in range(min(activation_map.shape[0], 3)): # Limit to 3 channels for RGB
channel_map = activation_map[c]
channel_map = (channel_map - channel_map.min()) / (channel_map.max() - channel_map.min())
rgb_activation_map[:, :, c] = (channel_map * 255).astype(np.uint8)
activation_image = Image.fromarray(rgb_activation_map, mode='RGB')
# Save activation maps
class_pred = predicted_classes[i].item()
parts = filenames[i].split('_')
plate = parts[0]
well = parts[1]
save_class_dir = os.path.join(save_dir, f'class_{class_pred}', str(plate), str(well))
os.makedirs(save_class_dir, exist_ok=True)
save_path = os.path.join(save_class_dir, f'{filenames[i]}')
if settings['save']:
activation_image.save(save_path)
img_paths.append(save_path)
if settings['save']:
activation_maps_to_database(img_paths, source_folder, settings)
if settings['correlation']:
df = calculate_activation_correlations(inputs, activation_maps, filenames, manders_thresholds=settings['manders_thresholds'])
if settings['plot']:
display(df)
if settings['save']:
activation_correlations_to_database(df, img_paths, source_folder, settings)
stop = time.time()
duration = stop - start
time_ls.append(duration)
files_processed = batch_idx * settings['batch_size']
files_to_process = len(data_loader) * settings['batch_size']
print_progress(files_processed, files_to_process, n_jobs=n_jobs, time_ls=time_ls, batch_size=settings['batch_size'], operation_type="Generating Activation Maps")
torch.cuda.empty_cache()
gc.collect()
print("Activation map generation complete.")
[docs]
def visualize_classes(model, dtype, class_names, **kwargs):
from .utils import class_visualization
for target_y in range(2): # Assuming binary classification
print(f"Visualizing class: {class_names[target_y]}")
visualization = class_visualization(target_y, model, dtype, **kwargs)
plt.imshow(visualization)
plt.title(f"Class {class_names[target_y]} Visualization")
plt.axis('off')
plt.show()
[docs]
def visualize_integrated_gradients(src, model_path, target_label_idx=0, image_size=224, channels=[1,2,3], normalize=True, save_integrated_grads=False, save_dir='integrated_grads'):
from .utils import IntegratedGradients, preprocess_image
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
model = torch.load(model_path)
model.to(device)
integrated_gradients = IntegratedGradients(model)
if save_integrated_grads and not os.path.exists(save_dir):
os.makedirs(save_dir)
images = []
filenames = []
for file in os.listdir(src):
if not file.endswith('.png'):
continue
image_path = os.path.join(src, file)
image, input_tensor = preprocess_image(image_path, normalize=normalize, image_size=image_size, channels=channels)
images.append(image)
filenames.append(file)
input_tensor = input_tensor.to(device)
integrated_grads = integrated_gradients.generate_integrated_gradients(input_tensor, target_label_idx)
integrated_grads = np.mean(integrated_grads, axis=1).squeeze()
fig, ax = plt.subplots(1, 3, figsize=(20, 5))
ax[0].imshow(image)
ax[0].axis('off')
ax[0].set_title("Original Image")
ax[1].imshow(integrated_grads, cmap='hot')
ax[1].axis('off')
ax[1].set_title("Integrated Gradients")
overlay = np.array(image)
overlay = overlay / overlay.max()
integrated_grads_rgb = np.stack([integrated_grads] * 3, axis=-1) # Convert saliency map to RGB
overlay = (overlay * 0.5 + integrated_grads_rgb * 0.5).clip(0, 1)
ax[2].imshow(overlay)
ax[2].axis('off')
ax[2].set_title("Overlay")
plt.show()
if save_integrated_grads:
os.makedirs(save_dir, exist_ok=True)
integrated_grads_image = Image.fromarray((integrated_grads * 255).astype(np.uint8))
integrated_grads_image.save(os.path.join(save_dir, f'integrated_grads_{file}'))
[docs]
class SmoothGrad:
def __init__(self, model, n_samples=50, stdev_spread=0.15):
[docs]
self.n_samples = n_samples
[docs]
self.stdev_spread = stdev_spread
[docs]
def compute_smooth_grad(self, input_tensor, target_class):
self.model.eval()
stdev = self.stdev_spread * (input_tensor.max() - input_tensor.min())
total_gradients = torch.zeros_like(input_tensor)
for i in range(self.n_samples):
noise = torch.normal(mean=0, std=stdev, size=input_tensor.shape).to(input_tensor.device)
noisy_input = input_tensor + noise
noisy_input.requires_grad_()
output = self.model(noisy_input)
self.model.zero_grad()
output[0, target_class].backward()
total_gradients += noisy_input.grad
avg_gradients = total_gradients / self.n_samples
return avg_gradients.abs()
[docs]
def visualize_smooth_grad(src, model_path, target_label_idx, image_size=224, channels=[1,2,3], normalize=True, save_smooth_grad=False, save_dir='smooth_grad'):
from .utils import preprocess_image
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
model = torch.load(model_path)
model.to(device)
smooth_grad = SmoothGrad(model)
if save_smooth_grad and not os.path.exists(save_dir):
os.makedirs(save_dir)
images = []
filenames = []
for file in os.listdir(src):
if not file.endswith('.png'):
continue
image_path = os.path.join(src, file)
image, input_tensor = preprocess_image(image_path, normalize=normalize, image_size=image_size, channels=channels)
images.append(image)
filenames.append(file)
input_tensor = input_tensor.to(device)
smooth_grad_map = smooth_grad.compute_smooth_grad(input_tensor, target_label_idx)
smooth_grad_map = np.mean(smooth_grad_map.cpu().data.numpy(), axis=1).squeeze()
fig, ax = plt.subplots(1, 3, figsize=(20, 5))
ax[0].imshow(image)
ax[0].axis('off')
ax[0].set_title("Original Image")
ax[1].imshow(smooth_grad_map, cmap='hot')
ax[1].axis('off')
ax[1].set_title("SmoothGrad")
overlay = np.array(image)
overlay = overlay / overlay.max()
smooth_grad_map_rgb = np.stack([smooth_grad_map] * 3, axis=-1) # Convert smooth grad map to RGB
overlay = (overlay * 0.5 + smooth_grad_map_rgb * 0.5).clip(0, 1)
ax[2].imshow(overlay)
ax[2].axis('off')
ax[2].set_title("Overlay")
plt.show()
if save_smooth_grad:
os.makedirs(save_dir, exist_ok=True)
smooth_grad_image = Image.fromarray((smooth_grad_map * 255).astype(np.uint8))
smooth_grad_image.save(os.path.join(save_dir, f'smooth_grad_{file}'))
[docs]
def deep_spacr(settings={}):
from .settings import deep_spacr_defaults
from .io import generate_training_dataset, generate_dataset
from .utils import save_settings
settings = deep_spacr_defaults(settings)
src = settings['src']
save_settings(settings, name='DL_model')
if settings['train'] or settings['test']:
if settings['generate_training_dataset']:
print(f"Generating train and test datasets ...")
train_path, test_path = generate_training_dataset(settings)
print(f'Generated Train set: {train_path}')
print(f'Generated Test set: {test_path}')
settings['src'] = os.path.dirname(train_path)
if settings['train_DL_model']:
print(f"Training model ...")
model_path = train_test_model(settings)
settings['model_path'] = model_path
settings['src'] = src
if settings['apply_model_to_dataset']:
if not settings['tar_path'] and os.path.isabs(settings['tar_path']) and os.path.exists(settings['tar_path']):
print(f"{settings['tar_path']} not found generating dataset ...")
tar_path = generate_dataset(settings)
settings['tar_path'] = tar_path
if os.path.exists(settings['model_path']):
apply_model_to_tar(settings)
[docs]
def model_knowledge_transfer(teacher_paths, student_save_path, data_loader, device='cpu', student_model_name='maxvit_t', pretrained=True, dropout_rate=None, use_checkpoint=False, alpha=0.5, temperature=2.0, lr=1e-4, epochs=10):
from spacr.utils import TorchModel # Adapt if needed
# Adjust filename to reflect knowledge-distillation if desired
if student_save_path.endswith('.pth'):
base, ext = os.path.splitext(student_save_path)
else:
base = student_save_path
student_save_path = base + '_KD.pth'
# -- 1. Load teacher models --
teachers = []
print("Loading teacher models:")
for path in teacher_paths:
print(f" Loading teacher: {path}")
ckpt = torch.load(path, map_location=device)
if isinstance(ckpt, TorchModel):
teacher = ckpt.to(device)
elif isinstance(ckpt, dict):
# If it's a dict with 'model' inside
# We might need to check if it has 'model_name', etc.
# But let's keep it simple: same architecture as the student
teacher = TorchModel(
model_name=ckpt.get('model_name', student_model_name),
pretrained=ckpt.get('pretrained', pretrained),
dropout_rate=ckpt.get('dropout_rate', dropout_rate),
use_checkpoint=ckpt.get('use_checkpoint', use_checkpoint)
).to(device)
teacher.load_state_dict(ckpt['model'])
else:
raise ValueError(f"Unsupported checkpoint type at {path} (must be TorchModel or dict).")
teacher.eval() # For consistent batchnorm, dropout
teachers.append(teacher)
# -- 2. Initialize the student TorchModel --
student_model = TorchModel(
model_name=student_model_name,
pretrained=pretrained,
dropout_rate=dropout_rate,
use_checkpoint=use_checkpoint
).to(device)
# You could load a partial checkpoint into the student here if desired.
# -- 3. Optimizer --
optimizer = optim.Adam(student_model.parameters(), lr=lr)
# Distillation training loop
for epoch in range(epochs):
student_model.train()
running_loss = 0.0
for images, labels in data_loader:
images, labels = images.to(device), labels.to(device)
# Forward pass student
logits_s = student_model(images) # shape: (B, num_classes)
logits_s_temp = logits_s / temperature # scale by T
# Distillation from teachers
with torch.no_grad():
# We'll average teacher probabilities
teacher_probs_list = []
for tm in teachers:
logits_t = tm(images) / temperature
# convert to probabilities
teacher_probs_list.append(F.softmax(logits_t, dim=1))
# average them
teacher_probs_ensemble = torch.mean(torch.stack(teacher_probs_list), dim=0)
# Student probabilities (log-softmax)
student_log_probs = F.log_softmax(logits_s_temp, dim=1)
# Distillation loss => KLDiv
loss_distill = F.kl_div(
student_log_probs,
teacher_probs_ensemble,
reduction='batchmean'
) * (temperature ** 2)
# Real label loss => cross-entropy
# We can compute this on the raw logits or scaled. Typically raw logits is standard:
loss_ce = F.cross_entropy(logits_s, labels)
# Weighted sum
loss = alpha * loss_ce + (1 - alpha) * loss_distill
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
avg_loss = running_loss / len(data_loader)
print(f"Epoch [{epoch+1}/{epochs}] - Loss: {avg_loss:.4f}")
# -- 4. Save final student as a TorchModel --
torch.save(student_model, student_save_path)
print(f"Knowledge-distilled student saved to: {student_save_path}")
return student_model
[docs]
def model_fusion(model_paths,save_path,device='cpu',model_name='maxvit_t',pretrained=True,dropout_rate=None,use_checkpoint=False,aggregator='mean'):
from spacr.utils import TorchModel
if save_path.endswith('.pth'):
save_path_part1, ext = os.path.splitext(save_path)
else:
save_path_part1 = save_path
save_path = save_path_part1 + f'_{aggregator}.pth'
valid_aggregators = {'mean', 'geomean', 'median', 'sum', 'max', 'min'}
if aggregator not in valid_aggregators:
raise ValueError(f"Invalid aggregator '{aggregator}'. "
f"Must be one of {valid_aggregators}.")
# --- 1. Load the first checkpoint to figure out architecture & hyperparams ---
print(f"Loading the first model from: {model_paths[0]} to derive architecture")
first_ckpt = torch.load(model_paths[0], map_location=device)
if isinstance(first_ckpt, dict):
# It's a dict with state_dict + possibly metadata
# Use any stored metadata if present
model_name = first_ckpt.get('model_name', model_name)
pretrained = first_ckpt.get('pretrained', pretrained)
dropout_rate = first_ckpt.get('dropout_rate', dropout_rate)
use_checkpoint = first_ckpt.get('use_checkpoint', use_checkpoint)
# Initialize the fused model
fused_model = TorchModel(
model_name=model_name,
pretrained=pretrained,
dropout_rate=dropout_rate,
use_checkpoint=use_checkpoint
).to(device)
# We'll collect state dicts in a list
state_dicts = [first_ckpt['model']] # the actual weights
elif isinstance(first_ckpt, TorchModel):
# The checkpoint is directly a TorchModel instance
fused_model = first_ckpt.to(device)
state_dicts = [fused_model.state_dict()]
else:
raise ValueError("Unsupported checkpoint format. Must be a dict or a TorchModel instance.")
# --- 2. Load the rest of the checkpoints ---
for path in model_paths[1:]:
print(f"Loading model from: {path}")
ckpt = torch.load(path, map_location=device)
if isinstance(ckpt, dict):
state_dicts.append(ckpt['model']) # Just the state dict portion
elif isinstance(ckpt, TorchModel):
state_dicts.append(ckpt.state_dict())
else:
raise ValueError(f"Unsupported checkpoint format in {path} (must be dict or TorchModel).")
# --- 3. Verify all state dicts have the same keys ---
fused_sd = fused_model.state_dict()
for sd in state_dicts:
if fused_sd.keys() != sd.keys():
raise ValueError("All models must have identical architecture/state_dict keys.")
# --- 4. Define aggregator logic ---
def combine_tensors(tensor_list, mode='mean'):
"""Given a list of Tensors, combine them using the chosen aggregator."""
# stack along new dimension => shape (num_models, *tensor.shape)
stacked = torch.stack(tensor_list, dim=0).float()
if mode == 'mean':
return stacked.mean(dim=0)
elif mode == 'geomean':
# geometric mean = exp(mean(log(tensor)))
# caution: requires all > 0
return torch.exp(torch.log(stacked).mean(dim=0))
elif mode == 'median':
return stacked.median(dim=0).values
elif mode == 'sum':
return stacked.sum(dim=0)
elif mode == 'max':
return stacked.max(dim=0).values
elif mode == 'min':
return stacked.min(dim=0).values
else:
raise ValueError(f"Unsupported aggregator: {mode}")
# --- 5. Combine the weights ---
for key in fused_sd.keys():
# gather all versions of this tensor
all_tensors = [sd[key] for sd in state_dicts]
fused_sd[key] = combine_tensors(all_tensors, mode=aggregator)
# Load combined weights into the fused model
fused_model.load_state_dict(fused_sd)
# --- 6. Save the entire TorchModel object ---
torch.save(fused_model, save_path)
print(f"Fused model (aggregator='{aggregator}') saved as a full TorchModel to: {save_path}")
return fused_model
[docs]
def annotate_filter_vision(settings):
from .utils import annotate_conditions, correct_metadata
def filter_csv_by_png(csv_file):
"""
Filters a DataFrame by removing rows that match PNG filenames in a folder.
Parameters:
csv_file (str): Path to the CSV file.
Returns:
pd.DataFrame: Filtered DataFrame.
"""
# Split the path to identify the datasets folder and build the training folder path
before_datasets, after_datasets = csv_file.split(os.sep + "datasets" + os.sep, 1)
train_fldr = os.path.join(before_datasets, 'datasets', 'training', 'train')
# Paths for train/nc and train/pc
nc_folder = os.path.join(train_fldr, 'nc')
pc_folder = os.path.join(train_fldr, 'pc')
# Load the CSV file into a DataFrame
df = pd.read_csv(csv_file)
# Collect PNG filenames from train/nc and train/pc
png_files = set()
for folder in [nc_folder, pc_folder]:
if os.path.exists(folder): # Ensure the folder exists
png_files.update({file for file in os.listdir(folder) if file.endswith(".png")})
# Filter the DataFrame by excluding rows where filenames match PNG files
filtered_df = df[~df['path'].isin(png_files)]
return filtered_df
if isinstance(settings['src'], str):
settings['src'] = [settings['src']]
for src in settings['src']:
ann_src, ext = os.path.splitext(src)
output_csv = ann_src+'_annotated_filtered.csv'
print(output_csv)
df = pd.read_csv(src)
df = correct_metadata(df)
df = annotate_conditions(df,
cells=settings['cells'],
cell_loc=settings['cell_loc'],
pathogens=settings['pathogens'],
pathogen_loc=settings['pathogen_loc'],
treatments=settings['treatments'],
treatment_loc=settings['treatment_loc'])
if not settings['filter_column'] is None:
if settings['filter_column'] in df.columns:
filtered_df = df[(df[settings['filter_column']] > settings['upper_threshold']) | (df[settings['filter_column']] < settings['lower_threshold'])]
print(f'Filtered DataFrame with {len(df)} rows to {len(filtered_df)} rows.')
else:
print(f"{settings['filter_column']} not in DataFrame columns.")
filtered_df = df
else:
filtered_df = df
filtered_df.to_csv(output_csv, index=False)
if settings['remove_train']:
df = filter_csv_by_png(output_csv)
df.to_csv(output_csv, index=False)