spacr 0.2.56__py3-none-any.whl → 0.2.61__py3-none-any.whl

spacr/gui_elements.py

@@ -516,6 +516,7 @@ class spacrDropdownMenu(tk.Frame):
         self.inactive_color = color_settings['inactive_color']
         self.active_color = color_settings['active_color']
         self.fg_color = color_settings['fg_color']
+        self.bg_color = style_out['bg_color']
 
         # Create the button with rounded edges
         self.button_bg = self.create_rounded_rectangle(2, 2, self.button_width + 2, self.size + 2, radius=20, fill=self.inactive_color, outline=self.inactive_color)
@@ -536,8 +537,8 @@ class spacrDropdownMenu(tk.Frame):
         self.canvas.bind("<Leave>", self.on_leave)
         self.canvas.bind("<Button-1>", self.on_click)
 
-        # Create a popup menu
-        self.menu = tk.Menu(self, tearoff=0)
+        # Create a popup menu with the desired background color
+        self.menu = tk.Menu(self, tearoff=0, bg=self.bg_color, fg=self.fg_color)
         for option in self.options:
             self.menu.add_command(label=option, command=lambda opt=option: self.on_select(opt))
 
@@ -591,7 +592,6 @@ class spacrDropdownMenu(tk.Frame):
                 else:
                     self.menu.entryconfig(idx, background=style_out['bg_color'], foreground=style_out['fg_color'])
 
-
 class spacrCheckbutton(ttk.Checkbutton):
     def __init__(self, parent, text="", variable=None, command=None, *args, **kwargs):
         super().__init__(parent, *args, **kwargs)
@@ -613,17 +613,26 @@ class spacrProgressBar(ttk.Progressbar):
         self.bg_color = style_out['bg_color']
         self.active_color = style_out['active_color']
         self.inactive_color = style_out['inactive_color']
+        self.font_size = style_out['font_size']
+        self.font_loader = style_out['font_loader']
 
         # Configure the style for the progress bar
         self.style = ttk.Style()
+        
+        # Remove any borders and ensure the active color fills the entire space
         self.style.configure(
             "spacr.Horizontal.TProgressbar",
-            troughcolor=self.bg_color,
-            background=self.active_color,
-            thickness=20,
-            troughrelief='flat',
-            borderwidth=0
+            troughcolor=self.inactive_color,       # Set the trough to bg color
+            background=self.active_color,    # Active part is the active color
+            borderwidth=0,                   # Remove border width
+            pbarrelief="flat",               # Flat relief for the progress bar
+            troughrelief="flat",             # Flat relief for the trough
+            thickness=20,                    # Set the thickness of the progress bar
+            darkcolor=self.active_color,     # Ensure darkcolor matches the active color
+            lightcolor=self.active_color,    # Ensure lightcolor matches the active color
+            bordercolor=self.bg_color        # Set the border color to the background color to hide it
         )
+
         self.configure(style="spacr.Horizontal.TProgressbar")
 
         # Set initial value to 0
@@ -641,15 +650,14 @@ class spacrProgressBar(ttk.Progressbar):
                 justify='left',
                 bg=self.inactive_color,
                 fg=self.fg_color,
-                wraplength=300  # Adjust the wraplength as needed
+                wraplength=300,
+                font=self.font_loader.get_font(size=self.font_size)
             )
-            self.progress_label.grid_forget()  # Temporarily hide it
+            self.progress_label.grid_forget()
 
         # Initialize attributes for time and operation
         self.operation_type = None
-        self.time_image = None
-        self.time_batch = None
-        self.time_left = None
+        self.additional_info = None
 
     def set_label_position(self):
         if self.label and self.progress_label:
@@ -664,74 +672,19 @@ class spacrProgressBar(ttk.Progressbar):
             label_text = f"Processing: {self['value']}/{self['maximum']}"
             if self.operation_type:
                 label_text += f", {self.operation_type}"
-            if self.time_image:
-                label_text += f", Time/image: {self.time_image:.3f} sec"
-            if self.time_batch:
-                label_text += f", Time/batch: {self.time_batch:.3f} sec"
-            if self.time_left:
-                label_text += f", Time_left: {self.time_left:.3f} min"
-            self.progress_label.config(text=label_text)
-
-class spacrProgressBar_v1(ttk.Progressbar):
-    def __init__(self, parent, label=True, *args, **kwargs):
-        super().__init__(parent, *args, **kwargs)
-
-        # Get the style colors
-        style_out = set_dark_style(ttk.Style())
-
-        self.fg_color = style_out['fg_color']
-        self.bg_color = style_out['bg_color']
-        self.active_color = style_out['active_color']
-        self.inactive_color = style_out['inactive_color']
-
-        # Configure the style for the progress bar
-        self.style = ttk.Style()
-        self.style.configure(
-            "spacr.Horizontal.TProgressbar",
-            troughcolor=self.bg_color,
-            background=self.active_color,
-            thickness=20,
-            troughrelief='flat',
-            borderwidth=0
-        )
-        self.configure(style="spacr.Horizontal.TProgressbar")
-
-        # Set initial value to 0
-        self['value'] = 0
-
-        # Track whether to show the progress label
-        self.label = label
-
-        # Create the progress label (defer placement)
-        if self.label:
-            self.progress_label = tk.Label(parent, text="Processing: 0/0", anchor='w', justify='left', bg=self.inactive_color, fg=self.fg_color)
-            self.progress_label.grid_forget()  # Temporarily hide it
-
-        # Initialize attributes for time and operation
-        self.operation_type = None
-        self.time_image = None
-        self.time_batch = None
-        self.time_left = None
-
-    def set_label_position(self):
-        if self.label and self.progress_label:
-            row_info = self.grid_info().get('row', 0)
-            col_info = self.grid_info().get('column', 0)
-            col_span = self.grid_info().get('columnspan', 1)
-            self.progress_label.grid(row=row_info + 1, column=col_info, columnspan=col_span, pady=5, padx=5, sticky='ew')
-
-    def update_label(self):
-        if self.label and self.progress_label:
-            # Update the progress label with current progress and additional info
-            label_text = f"Processing: {self['value']}/{self['maximum']}"
-            if self.operation_type:
-                label_text += f", {self.operation_type}"
-            if self.time_image:
-                label_text += f", Time/image: {self.time_image:.3f} sec"
-            if self.time_batch:
-                label_text += f", Time/batch: {self.time_batch:.3f} sec"
-            if self.time_left:
-                label_text += f", Time_left: {self.time_left:.3f} min"
+            if hasattr(self, 'additional_info') and self.additional_info:
+                # Add a space between progress information and additional information
+                label_text += "\n\n"
+                # Split the additional_info into a list of items
+                items = self.additional_info.split(", ")
+                formatted_additional_info = ""
+                # Group the items in pairs, adding them to formatted_additional_info
+                for i in range(0, len(items), 2):
+                    if i + 1 < len(items):
+                        formatted_additional_info += f"{items[i]}, {items[i + 1]}\n\n"
+                    else:
+                        formatted_additional_info += f"{items[i]}\n\n"  # If there's an odd item out, add it alone
+                label_text += formatted_additional_info.strip()
             self.progress_label.config(text=label_text)
 
 def spacrScrollbarStyle(style, inactive_color, active_color):

spacr/gui_utils.py

@@ -1,10 +1,11 @@
-import os, io, sys, ast, ctypes, ast, sqlite3, requests, time, traceback
+import os, io, sys, ast, ctypes, ast, sqlite3, requests, time, traceback, torch
 import tkinter as tk
 from tkinter import ttk
 import matplotlib
 import matplotlib.pyplot as plt
 matplotlib.use('Agg')
 from huggingface_hub import list_repo_files
+import psutil
 
 from .gui_elements import AnnotateApp, spacrEntry, spacrCheck, spacrCombo
 
@@ -12,6 +13,36 @@ try:
     ctypes.windll.shcore.SetProcessDpiAwareness(True)
 except AttributeError:
     pass
+
+def initialize_cuda():
+    """
+    Initializes CUDA in the main process by performing a simple GPU operation.
+    """
+    if torch.cuda.is_available():
+        # Allocate a small tensor on the GPU
+        _ = torch.tensor([0.0], device='cuda')
+        print("CUDA initialized in the main process.")
+    else:
+        print("CUDA is not available.")
+
+def set_high_priority(process):
+    try:
+        p = psutil.Process(process.pid)
+        if os.name == 'nt':  # Windows
+            p.nice(psutil.HIGH_PRIORITY_CLASS)
+        else:  # Unix-like systems
+            p.nice(-10)  # Adjusted priority level
+        print(f"Successfully set high priority for process: {process.pid}")
+    except psutil.AccessDenied as e:
+        print(f"Access denied when trying to set high priority for process {process.pid}: {e}")
+    except psutil.NoSuchProcess as e:
+        print(f"No such process {process.pid}: {e}")
+    except Exception as e:
+        print(f"Failed to set high priority for process {process.pid}: {e}")
+
+def set_cpu_affinity(process):
+    p = psutil.Process(process.pid)
+    p.cpu_affinity(list(range(os.cpu_count())))
     
 def proceed_with_app(root, app_name, app_func):
     # Clear the current content frame
@@ -44,16 +75,6 @@ def load_app(root, app_name, app_func):
     else:
         proceed_with_app(root, app_name, app_func)
     
-def parse_list_v1(value):
-    try:
-        parsed_value = ast.literal_eval(value)
-        if isinstance(parsed_value, list):
-            return parsed_value
-        else:
-            raise ValueError(f"Expected a list but got {type(parsed_value).__name__}")
-    except (ValueError, SyntaxError) as e:
-        raise ValueError(f"Invalid format for list: {value}. Error: {e}")
-    
 def parse_list(value):
     try:
         parsed_value = ast.literal_eval(value)
@@ -83,6 +104,7 @@ def create_input_field(frame, label_text, row, var_type='entry', options=None, d
     size_dict['settings_width'] = size_dict['settings_width'] - int(size_dict['settings_width']*0.1)
 
     # Replace underscores with spaces and capitalize the first letter
+
     label_text = label_text.replace('_', ' ').capitalize()
 
     # Configure the column widths
@@ -97,32 +119,35 @@ def create_input_field(frame, label_text, row, var_type='entry', options=None, d
     custom_frame.config(highlightbackground=style_out['bg_color'], highlightthickness=1, bd=2)
 
     # Create and configure the label
-    if font_loader:
-        label = ttk.Label(custom_frame, text=label_text, background=style_out['bg_color'], foreground=style_out['fg_color'], font=font_loader.get_font(size=font_size), anchor='e', justify='right')
-    label = ttk.Label(custom_frame, text=label_text, background=style_out['bg_color'], foreground=style_out['fg_color'], font=(style_out['font_family'], style_out['font_size']), anchor='e', justify='right')
+    label = tk.Label(custom_frame, text=label_text, bg=style_out['bg_color'], fg=style_out['fg_color'], font=font_loader.get_font(size=font_size), anchor='e', justify='right')
     label.grid(column=label_column, row=0, sticky=tk.W, padx=(5, 2), pady=5)  # Place the label in the first row
 
     # Create and configure the input widget based on var_type
-    if var_type == 'entry':
-        var = tk.StringVar(value=default_value)
-        entry = spacrEntry(custom_frame, textvariable=var, outline=False, width=size_dict['settings_width'])
-        entry.grid(column=widget_column, row=1, sticky=tk.W, padx=(2, 5), pady=5)  # Place the entry in the second row
-        return (label, entry, var, custom_frame)  # Return both the label and the entry, and the variable
-    elif var_type == 'check':
-        var = tk.BooleanVar(value=default_value)  # Set default value (True/False)
-        check = spacrCheck(custom_frame, text="", variable=var)
-        check.grid(column=widget_column, row=1, sticky=tk.W, padx=(2, 5), pady=5)  # Place the checkbutton in the second row
-        return (label, check, var, custom_frame)  # Return both the label and the checkbutton, and the variable
-    elif var_type == 'combo':
-        var = tk.StringVar(value=default_value)  # Set default value
-        combo = spacrCombo(custom_frame, textvariable=var, values=options, width=size_dict['settings_width'])  # Apply TCombobox style
-        combo.grid(column=widget_column, row=1, sticky=tk.W, padx=(2, 5), pady=5)  # Place the combobox in the second row
-        if default_value:
-            combo.set(default_value)
-        return (label, combo, var, custom_frame)  # Return both the label and the combobox, and the variable
-    else:
-        var = None  # Placeholder in case of an undefined var_type
-        return (label, None, var, custom_frame)
+    try:
+        if var_type == 'entry':
+            var = tk.StringVar(value=default_value)
+            entry = spacrEntry(custom_frame, textvariable=var, outline=False, width=size_dict['settings_width'])
+            entry.grid(column=widget_column, row=1, sticky=tk.W, padx=(2, 5), pady=5)  # Place the entry in the second row
+            return (label, entry, var, custom_frame)  # Return both the label and the entry, and the variable
+        elif var_type == 'check':
+            var = tk.BooleanVar(value=default_value)  # Set default value (True/False)
+            check = spacrCheck(custom_frame, text="", variable=var)
+            check.grid(column=widget_column, row=1, sticky=tk.W, padx=(2, 5), pady=5)  # Place the checkbutton in the second row
+            return (label, check, var, custom_frame)  # Return both the label and the checkbutton, and the variable
+        elif var_type == 'combo':
+            var = tk.StringVar(value=default_value)  # Set default value
+            combo = spacrCombo(custom_frame, textvariable=var, values=options, width=size_dict['settings_width'])  # Apply TCombobox style
+            combo.grid(column=widget_column, row=1, sticky=tk.W, padx=(2, 5), pady=5)  # Place the combobox in the second row
+            if default_value:
+                combo.set(default_value)
+            return (label, combo, var, custom_frame)  # Return both the label and the combobox, and the variable
+        else:
+            var = None  # Placeholder in case of an undefined var_type
+            return (label, None, var, custom_frame)
+    except Exception as e:
+        traceback.print_exc()
+        print(f"Error creating input field: {e}")
+        print(f"Wrong type for {label_text} Expected {var_type}")
 
 def process_stdout_stderr(q):
     """
@@ -150,16 +175,6 @@ def cancel_after_tasks(frame):
             frame.after_cancel(task)
         frame.after_tasks.clear()
 
-def main_thread_update_function(root, q, fig_queue, canvas_widget):
-    try:
-        #ansi_escape_pattern = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')
-        while not q.empty():
-            message = q.get_nowait()
-    except Exception as e:
-        print(f"Error updating GUI canvas: {e}")
-    finally:
-        root.after(100, lambda: main_thread_update_function(root, q, fig_queue, canvas_widget))
-
 def annotate(settings):
     from .settings import set_annotate_default_settings
     settings = set_annotate_default_settings(settings)
@@ -341,7 +356,7 @@ def convert_settings_dict_for_gui(settings):
         'channels': ('combo', ['[0,1,2,3]', '[0,1,2]', '[0,1]', '[0]'], '[0,1,2,3]'),
         'train_channels': ('combo', ["['r','g','b']", "['r','g']", "['r','b']", "['g','b']", "['r']", "['g']", "['b']"], "['r','g','b']"),
         'channel_dims': ('combo', ['[0,1,2,3]', '[0,1,2]', '[0,1]', '[0]'], '[0,1,2,3]'),
-        'dataset_mode': ('combo', ['annotation', 'metadata', 'recruitment'], 'annotate'),
+        'dataset_mode': ('combo', ['annotation', 'metadata', 'recruitment'], 'metadata'),
         'cell_mask_dim': ('combo', chans, None),
         'cell_chann_dim': ('combo', chans, None),
         'nucleus_mask_dim': ('combo', chans, None),
@@ -476,7 +491,7 @@ def run_function_gui(settings_type, settings, q, fig_queue, stop_requested):
         imports = 2
     elif settings_type == 'recruitment':
         function = analyze_recruitment
-        imports = 2
+        imports = 1
     elif settings_type == 'umap':
         function = generate_image_umap
         imports = 1
@@ -490,7 +505,6 @@ def run_function_gui(settings_type, settings, q, fig_queue, stop_requested):
     finally:
         stop_requested.value = 1
 
-
 def hide_all_settings(vars_dict, categories):
     """
     Function to initially hide all settings in the GUI.

spacr/io.py

@@ -1,9 +1,9 @@
-import os, re, sqlite3, gc, torch, time, random, shutil, cv2, tarfile, cellpose, glob
+import os, re, sqlite3, gc, torch, time, random, shutil, cv2, tarfile, cellpose, glob, queue
 import numpy as np
 import pandas as pd
 import tifffile
-from PIL import Image
-from collections import defaultdict, Counter
+from PIL import Image, ImageOps
+from collections import defaultdict, Counter, deque
 from pathlib import Path
 from functools import partial
 from matplotlib.animation import FuncAnimation
@@ -17,12 +17,12 @@ import imageio.v2 as imageio2
 import matplotlib.pyplot as plt
 from io import BytesIO
 from IPython.display import display, clear_output
-from multiprocessing import Pool, cpu_count
-from torch.utils.data import Dataset
+from multiprocessing import Pool, cpu_count, Process, Queue
+from torch.utils.data import Dataset, DataLoader
 import matplotlib.pyplot as plt
 from torchvision.transforms import ToTensor
 import seaborn as sns
-
+import atexit
 
 from .logger import log_function_call
 
@@ -444,20 +444,7 @@ class NoClassDataset(Dataset):
         # Return both the image and its filename
         return img, self.filenames[index]
 
-class MyDataset(Dataset):
-    """
-    A custom dataset class for loading and processing image data.
-
-    Args:
-        data_dir (str): The directory path where the image data is stored.
-        loader_classes (list): A list of class names for the dataset.
-        transform (callable, optional): A function/transform to apply to the image data. Default is None.
-        shuffle (bool, optional): Whether to shuffle the dataset. Default is True.
-        pin_memory (bool, optional): Whether to pin the loaded images to memory. Default is False.
-        specific_files (list, optional): A list of specific file paths to include in the dataset. Default is None.
-        specific_labels (list, optional): A list of specific labels corresponding to the specific files. Default is None.
-    """
-
+class spacrDataset(Dataset):
     def __init__(self, data_dir, loader_classes, transform=None, shuffle=True, pin_memory=False, specific_files=None, specific_labels=None):
         self.data_dir = data_dir
         self.classes = loader_classes
@@ -466,7 +453,7 @@ class MyDataset(Dataset):
         self.pin_memory = pin_memory
         self.filenames = []
         self.labels = []
-        
+
         if specific_files and specific_labels:
             self.filenames = specific_files
             self.labels = specific_labels
@@ -479,33 +466,113 @@ class MyDataset(Dataset):
         
         if self.shuffle:
             self.shuffle_dataset()
-            
+
         if self.pin_memory:
-            self.images = [self.load_image(f) for f in self.filenames]
-    
+            # Use multiprocessing to load images in parallel
+            with Pool(processes=cpu_count()) as pool:
+                self.images = pool.map(self.load_image, self.filenames)
+        else:
+            self.images = None
+
     def load_image(self, img_path):
         img = Image.open(img_path).convert('RGB')
+        img = ImageOps.exif_transpose(img)  # Handle image orientation
         return img
-    
+
     def __len__(self):
         return len(self.filenames)
-    
+
     def shuffle_dataset(self):
         combined = list(zip(self.filenames, self.labels))
         random.shuffle(combined)
         self.filenames, self.labels = zip(*combined)
-        
+
     def get_plate(self, filepath):
-        filename = os.path.basename(filepath)  # Get just the filename from the full path
+        filename = os.path.basename(filepath)
         return filename.split('_')[0]
-    
+
     def __getitem__(self, index):
+        if self.pin_memory:
+            img = self.images[index]
+        else:
+            img = self.load_image(self.filenames[index])
         label = self.labels[index]
         filename = self.filenames[index]
-        img = self.load_image(filename)
         if self.transform:
             img = self.transform(img)
         return img, label, filename
+    
+class spacrDataLoader(DataLoader):
+    def __init__(self, *args, preload_batches=1, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.preload_batches = preload_batches
+        self.batch_queue = Queue(maxsize=preload_batches)
+        self.process = None
+        self.current_batch_index = 0
+        self._stop_event = False
+        self.pin_memory = kwargs.get('pin_memory', False)
+        atexit.register(self.cleanup)
+
+    def _preload_next_batches(self):
+        try:
+            for _ in range(self.preload_batches):
+                if self._stop_event:
+                    break
+                batch = next(self._iterator)
+                if self.pin_memory:
+                    batch = self._pin_memory_batch(batch)
+                self.batch_queue.put(batch)
+        except StopIteration:
+            pass
+
+    def _start_preloading(self):
+        if self.process is None or not self.process.is_alive():
+            self._iterator = iter(super().__iter__())
+            if not self.pin_memory:
+                self.process = Process(target=self._preload_next_batches)
+                self.process.start()
+            else:
+                self._preload_next_batches()  # Directly load if pin_memory is True
+
+    def _pin_memory_batch(self, batch):
+        if isinstance(batch, (list, tuple)):
+            return [b.pin_memory() if isinstance(b, torch.Tensor) else b for b in batch]
+        elif isinstance(batch, torch.Tensor):
+            return batch.pin_memory()
+        else:
+            return batch
+
+    def __iter__(self):
+        self._start_preloading()
+        return self
+
+    def __next__(self):
+        if self.process and not self.process.is_alive() and self.batch_queue.empty():
+            raise StopIteration
+
+        try:
+            if self.pin_memory:
+                next_batch = self.batch_queue.get(timeout=60)
+            else:
+                next_batch = self.batch_queue.get(timeout=60)
+            self.current_batch_index += 1
+
+            # Start preloading the next batches
+            if self.batch_queue.qsize() < self.preload_batches:
+                self._start_preloading()
+
+            return next_batch
+        except queue.Empty:
+            raise StopIteration
+
+    def cleanup(self):
+        self._stop_event = True
+        if self.process and self.process.is_alive():
+            self.process.terminate()
+            self.process.join()
+
+    def __del__(self):
+        self.cleanup()
 
 class NoClassDataset(Dataset):
     def __init__(self, data_dir, transform=None, shuffle=True, load_to_memory=False):
@@ -2292,7 +2359,7 @@ def _save_model(model, model_type, results_df, dst, epoch, epochs, intermedeate_
     
     def save_model_at_threshold(threshold, epoch, suffix=""):
         percentile = str(threshold * 100)
-        print(f'\rfound: {percentile}% accurate model')#, end='\r', flush=True)
+        print(f'Found: {percentile}% accurate model')
         model_path = f'{dst}/{model_type}_epoch_{str(epoch)}{suffix}_acc_{percentile}_channels_{channels_str}.pth'
         torch.save(model, model_path)
         return model_path
@@ -2303,7 +2370,8 @@ def _save_model(model, model_type, results_df, dst, epoch, epochs, intermedeate_
         return model_path
 
     for threshold in intermedeate_save:
-        if results_df['neg_accuracy'].dropna().mean() >= threshold and results_df['pos_accuracy'].dropna().mean() >= threshold:
+        if results_df['neg_accuracy'] >= threshold and results_df['pos_accuracy'] >= threshold:
+            print(f"Nc class accuracy: {results_df['neg_accuracy']} Pc class Accuracy: {results_df['pos_accuracy']}")
             model_path = save_model_at_threshold(threshold, epoch)
             break
         else:
@@ -2311,7 +2379,7 @@ def _save_model(model, model_type, results_df, dst, epoch, epochs, intermedeate_
     
     return model_path
 
-def _save_progress(dst, results_df, train_metrics_df, epoch, epochs):
+def _save_progress(dst, results_df, result_type='train'):
     """
     Save the progress of the classification model.
 
@@ -2325,18 +2393,13 @@ def _save_progress(dst, results_df, train_metrics_df, epoch, epochs):
     """
     # Save accuracy, loss, PRAUC
     os.makedirs(dst, exist_ok=True)
-    results_path = os.path.join(dst, 'acc_loss_prauc.csv')
+    results_path = os.path.join(dst, f'{result_type}.csv')
     if not os.path.exists(results_path):
         results_df.to_csv(results_path, index=True, header=True, mode='w')
     else:
         results_df.to_csv(results_path, index=True, header=False, mode='a')
-    
-    training_metrics_path = os.path.join(dst, 'training_metrics.csv')
-    if not os.path.exists(training_metrics_path):
-        train_metrics_df.to_csv(training_metrics_path, index=True, header=True, mode='w')
-    else:
-        train_metrics_df.to_csv(training_metrics_path, index=True, header=False, mode='a')
-    if epoch == epochs:
+
+    if result_type == 'train':
         read_plot_model_stats(results_path, save=True)
     return
 

spacr/plot.py

@@ -1,4 +1,4 @@
-import os,re, random, cv2, glob, time, math
+import os,re, random, cv2, glob, time, math, torch
 
 import numpy as np
 import pandas as pd
@@ -1186,6 +1186,52 @@ def _imshow(img, labels, nrow=20, color='white', fontsize=12):
         y = row * img_height + 15
         plt.text(x, y, label, color=color, fontsize=fontsize, fontweight='bold')
     return fig
+
+def _imshow_gpu(img, labels, nrow=20, color='white', fontsize=12):
+    """
+    Display multiple images in a grid with corresponding labels.
+
+    Args:
+        img (torch.Tensor): A batch of images as a tensor.
+        labels (list): List of labels corresponding to each image.
+        nrow (int, optional): Number of images per row in the grid. Defaults to 20.
+        color (str, optional): Color of the label text. Defaults to 'white'.
+        fontsize (int, optional): Font size of the label text. Defaults to 12.
+    """
+    if img.is_cuda:
+        img = img.cpu()  # Move to CPU if the tensor is on GPU
+
+    n_images = len(labels)
+    n_col = nrow
+    n_row = int(np.ceil(n_images / n_col))
+
+    img_height = img.shape[2]  # Height of the image
+    img_width = img.shape[3]   # Width of the image
+
+    # Prepare the canvas on CPU
+    canvas = torch.zeros((img_height * n_row, img_width * n_col, 3))
+
+    for i in range(n_row):
+        for j in range(n_col):
+            idx = i * n_col + j
+            if idx < n_images:
+                # Place the image on the canvas
+                canvas[i * img_height:(i + 1) * img_height, j * img_width:(j + 1) * img_width] = img[idx].permute(1, 2, 0)
+
+    canvas = canvas.numpy()  # Convert to NumPy for plotting
+
+    fig = plt.figure(figsize=(50, 50))
+    plt.imshow(canvas)
+    plt.axis("off")
+
+    for i, label in enumerate(labels):
+        row = i // n_col
+        col = i % n_col
+        x = col * img_width + 2
+        y = row * img_height + 15
+        plt.text(x, y, label, color=color, fontsize=fontsize, fontweight='bold')
+
+    return fig
     
 def _plot_histograms_and_stats(df):
     conditions = df['condition'].unique()