spacr 0.3.1__py3-none-any.whl → 0.3.22__py3-none-any.whl

spacr/deep_spacr.py

@@ -1,5 +1,4 @@
-import os, torch, time, gc, datetime
-
+import os, torch, time, gc, datetime, cv2
 torch.backends.cudnn.benchmark = True
 
 import numpy as np
@@ -8,15 +7,146 @@ from torch.optim import Adagrad, AdamW
 from torch.autograd import grad
 from torch.optim.lr_scheduler import StepLR
 import torch.nn.functional as F
-from IPython.display import display, clear_output
 import matplotlib.pyplot as plt
 from PIL import Image
 from sklearn.metrics import auc, precision_recall_curve
-from multiprocessing import set_start_method
-#set_start_method('spawn', force=True)
+from IPython.display import display
+from multiprocessing import cpu_count
+
+from torchvision import transforms
+from torch.utils.data import DataLoader
 
-from .logger import log_function_call
-from .utils import close_multiprocessing_processes, reset_mp
+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
+
+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'])
+    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]  
+    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
 
 def evaluate_model_performance(model, loader, epoch, loss_type):
     """
@@ -118,7 +248,7 @@ def evaluate_model_performance(model, loader, epoch, loss_type):
     
     loss /= len(loader)
     data_dict = classification_metrics(all_labels, prediction_pos_probs)
-    data_dict['loss'] = loss
+    data_dict['loss'] = loss.item()
     data_dict['epoch'] = epoch
     data_dict['Accuracy'] = acc
     
@@ -175,7 +305,7 @@ def test_model_core(model, loader, loader_name, epoch, loss_type):
         'class_1_probability':prediction_pos_probs})
 
     loss /= len(loader)
-    data_df = classification_metrics(all_labels, prediction_pos_probs, loader_name, loss, epoch)
+    data_df = classification_metrics(all_labels, prediction_pos_probs, loss, epoch)
     return data_df, prediction_pos_probs, all_labels, results_df
 
 def test_model_performance(loaders, model, loader_name_list, epoch, loss_type):
@@ -201,9 +331,12 @@ def test_model_performance(loaders, model, loader_name_list, epoch, loss_type):
 
 def train_test_model(settings):
     
-    from .io import _save_settings, _copy_missclassified
-    from .utils import pick_best_model
-    from .core import generate_loaders
+    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()
@@ -221,7 +354,12 @@ def train_test_model(settings):
         model = torch.load(settings['custom_model_path'])
     
     if settings['train']:
-        _save_settings(settings, src)
+        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, 
@@ -235,7 +373,6 @@ def train_test_model(settings):
                                                   normalize=settings['normalize'],
                                                   channels=settings['train_channels'],
                                                   augment=settings['augment'],
-                                                  preload_batches=settings['preload_batches'],
                                                   verbose=settings['verbose'])
         
         #train_batch_1_figure = os.path.join(dst, 'batch_1.pdf')
@@ -264,19 +401,19 @@ def train_test_model(settings):
                                         channels=settings['train_channels'])
         
     if settings['test']:
-        test, _, plate_names_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,
-                                                                preload_batches=settings['preload_batches'],
-                                                                verbose=settings['verbose'])
+        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}')
@@ -308,7 +445,10 @@ def train_test_model(settings):
     torch.cuda.memory.empty_cache()
     gc.collect()
 
-    return model_path
+    if settings['train']:
+        return model_path
+    if settings['test']:
+        return result_loc
     
 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):
     """
@@ -355,11 +495,6 @@ def train_model(dst, model_type, train_loaders, epochs=100, learning_rate=0.0001
     
     kwargs = {'n_jobs': n_jobs, 'pin_memory': True} if use_cuda else {}
     
-    #for idx, (images, labels, filenames) in enumerate(train_loaders):
-    #    batch, chans, height, width = images.shape
-    #    break
-
-    
     model = choose_model(model_type, device, init_weights, dropout_rate, use_checkpoint, verbose=verbose)
     
     
@@ -452,19 +587,21 @@ def train_model(dst, model_type, train_loaders, epochs=100, learning_rate=0.0001
             if schedule == 'step_lr':
                 scheduler.step()
 
-        if epoch % 10 == 0 or epoch == epochs:
-            if accumulated_train_dicts:
-                train_df = pd.DataFrame(accumulated_train_dicts)
-                _save_progress(dst, train_df, result_type='train')
-                
-            if accumulated_val_dicts:
-                val_df = pd.DataFrame(accumulated_val_dicts)
-                _save_progress(dst, val_df,result_type='validation')
-                
-            if accumulated_test_dicts:
-                val_df = pd.DataFrame(accumulated_test_dicts)
-                _save_progress(dst, val_df, result_type='test')
-                
+        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)
@@ -473,7 +610,138 @@ def train_model(dst, model_type, train_loaders, epochs=100, learning_rate=0.0001
             
     return model, model_path
 
-def visualize_saliency_map(src, model_type='maxvit', model_path='', image_size=224, channels=[1,2,3], normalize=True, class_names=None, save_saliency=False, save_dir='saliency_maps'):
+def visualize_saliency_map(settings):
+    from spacr.utils import SaliencyMapGenerator, print_progress
+    from spacr.io import TarImageDataset  # Assuming you have a dataset class
+    from torchvision.utils import make_grid
+
+    use_cuda = torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    # Set number of jobs for loading
+    if settings['n_jobs'] is None:
+        n_jobs = max(1, cpu_count() - 4)
+    else:
+        n_jobs = settings['n_jobs']
+
+    # Set transforms for images
+    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']))])
+
+    # Handle dataset path
+    if os.path.exists(settings['dataset']):
+        tar_path = settings['dataset']
+    else:
+        print(f"Dataset not found at {settings['dataset']}")
+        return
+    
+    if settings.get('save', False):
+        if settings['dtype'] not in ['uint8', 'uint16']:
+            print("Invalid dtype in settings. Please use 'uint8' or 'uint16'.")
+            return
+
+    # Load the model
+    model = torch.load(settings['model_path'])
+    model.to(device)
+    model.eval()  # Ensure the model is in evaluation mode
+
+    # Create directory for saving saliency maps if it does not exist
+    if settings.get('save', False):
+        dataset_dir = os.path.dirname(tar_path)
+        dataset_name = os.path.splitext(os.path.basename(tar_path))[0]
+        save_dir = os.path.join(dataset_dir, dataset_name, 'saliency_maps')
+        os.makedirs(save_dir, exist_ok=True)
+        print(f"Saliency maps will be saved in: {save_dir}")
+
+    # Load dataset
+    dataset = TarImageDataset(tar_path, transform=transform)
+    data_loader = DataLoader(dataset, batch_size=settings['batch_size'], shuffle=True, num_workers=n_jobs, pin_memory=True)
+
+    # Initialize SaliencyMapGenerator
+    cam_generator = SaliencyMapGenerator(model)
+    time_ls = []
+
+    for batch_idx, (inputs, filenames) in enumerate(data_loader):
+        start = time.time()
+        inputs = inputs.to(device)
+        
+        saliency_maps, predicted_classes = cam_generator.compute_saliency_and_predictions(inputs)
+
+        if settings['saliency_mode'] not in ['mean', 'sum']:
+            print("To generate channel average or sum saliency maps set saliency_mode to 'mean' or 'sum', respectively.")
+
+        if settings['saliency_mode'] == 'mean':
+            saliency_maps = saliency_maps.mean(dim=1, keepdim=True)
+
+        elif settings['saliency_mode'] == 'sum':
+            saliency_maps = saliency_maps.sum(dim=1, keepdim=True)
+
+        # Example usage with the class
+        if settings.get('plot', False):
+            if settings['plot_mode'] not in ['mean', 'channel', '3-channel']:
+                print("Invalid plot_mode in settings. Please use 'mean', 'channel', or '3-channel'.")
+                return
+            else:
+                cam_generator.plot_saliency_grid(inputs, saliency_maps, predicted_classes, mode=settings['plot_mode'])
+
+        if settings.get('save', False):
+            for i in range(inputs.size(0)):
+                saliency_map = saliency_maps[i].detach().cpu().numpy()
+
+                # Check dtype in settings and normalize accordingly
+                if settings['dtype'] == 'uint16':
+                    saliency_map = np.clip(saliency_map, 0, 1) * 65535
+                    saliency_map = saliency_map.astype(np.uint16)
+                    mode = 'I;16'
+                elif settings['dtype'] == 'uint8':
+                    saliency_map = np.clip(saliency_map, 0, 1) * 255
+                    saliency_map = saliency_map.astype(np.uint8)
+                    mode = 'L'  # Grayscale mode for uint8
+
+                # Get the class prediction (0 or 1)
+                class_pred = predicted_classes[i].item()
+
+                save_class_dir = os.path.join(save_dir, f'class_{class_pred}')
+                os.makedirs(save_class_dir, exist_ok=True)
+                save_path = os.path.join(save_class_dir, filenames[i])
+
+                # Handle different cases based on saliency_map dimensions
+                if saliency_map.ndim == 3:  # Multi-channel case (C, H, W)
+                    if saliency_map.shape[0] == 3:  # RGB-like saliency map
+                        saliency_image = Image.fromarray(np.moveaxis(saliency_map, 0, -1), mode="RGB")  # Convert (C, H, W) to (H, W, C)
+                    elif saliency_map.shape[0] == 1:  # Single-channel case (1, H, W)
+                        saliency_map = np.squeeze(saliency_map)  # Remove the extra channel dimension
+                        saliency_image = Image.fromarray(saliency_map, mode=mode)  # Use grayscale mode for single-channel
+                    else:
+                        raise ValueError(f"Unexpected number of channels: {saliency_map.shape[0]}")
+
+                elif saliency_map.ndim == 2:  # Single-channel case (H, W)
+                    saliency_image = Image.fromarray(saliency_map, mode=mode)  # Keep single channel (H, W)
+
+                else:
+                    raise ValueError(f"Unexpected number of dimensions: {saliency_map.ndim}")
+
+                # Save the image
+                saliency_image.save(save_path)
+
+
+        stop = time.time()
+        duration = stop - start
+        time_ls.append(duration)
+        files_processed = batch_idx * settings['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=settings['batch_size'], operation_type="Generating Saliency Maps")
+
+    print("Saliency map generation complete.")
+
+def visualize_saliency_map_v1(src, model_type='maxvit', model_path='', image_size=224, channels=[1,2,3], normalize=True, class_names=None, save_saliency=False, save_dir='saliency_maps'):
 
     from spacr.utils import SaliencyMapGenerator, preprocess_image
 
@@ -732,7 +1000,7 @@ def visualize_smooth_grad(src, model_path, target_label_idx, image_size=224, cha
 
 def deep_spacr(settings={}):
     from .settings import deep_spacr_defaults
-    from .core import generate_training_dataset, generate_dataset, apply_model_to_tar
+    from .io import generate_training_dataset, generate_dataset
     from .utils import save_settings
     
     settings = deep_spacr_defaults(settings)

spacr/gui.py

@@ -48,6 +48,7 @@ class MainApp(tk.Tk):
         }
 
         self.additional_gui_apps = {
+            "Convert": (lambda frame: initiate_root(self, 'convert'), "Convert images to Grayscale TIFs."),
             "Umap": (lambda frame: initiate_root(self, 'umap'), "Generate UMAP embeddings with datapoints represented as images."),
             "Train Cellpose": (lambda frame: initiate_root(self, 'train_cellpose'), "Train custom Cellpose models."),
             "ML Analyze": (lambda frame: initiate_root(self, 'ml_analyze'), "Machine learning analysis of data."),