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

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,17 +75,19 @@ 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):
+    """
+    Parses a string representation of a list and returns the parsed list.
+
+    Args:
+        value (str): The string representation of the list.
+
+    Returns:
+        list: The parsed list.
+
+    Raises:
+        ValueError: If the input value is not a valid list format or contains mixed types or unsupported types.
+    """
     try:
         parsed_value = ast.literal_eval(value)
         if isinstance(parsed_value, list):
@@ -72,7 +105,26 @@ def parse_list(value):
 
 # Usage example in your create_input_field function
 def create_input_field(frame, label_text, row, var_type='entry', options=None, default_value=None):
+    """
+    Create an input field in the specified frame.
+
+    Args:
+        frame (tk.Frame): The frame in which the input field will be created.
+        label_text (str): The text to be displayed as the label for the input field.
+        row (int): The row in which the input field will be placed.
+        var_type (str, optional): The type of input field to create. Defaults to 'entry'.
+        options (list, optional): The list of options for a combo box input field. Defaults to None.
+        default_value (str, optional): The default value for the input field. Defaults to None.
+
+    Returns:
+        tuple: A tuple containing the label, input widget, variable, and custom frame.
+
+    Raises:
+        Exception: If an error occurs while creating the input field.
+
+    """
     from .gui_elements import set_dark_style, set_element_size
+    
     label_column = 0
     widget_column = 0  # Both label and widget will be in the same column
 
@@ -83,6 +135,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 +150,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 +206,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)
@@ -190,6 +236,12 @@ def annotate(settings):
 
 def generate_annotate_fields(frame):
     from .settings import set_annotate_default_settings
+    from .gui_elements import set_dark_style
+
+    style_out = set_dark_style(ttk.Style())
+    font_loader = style_out['font_loader']
+    font_size = style_out['font_size'] - 2
+
     vars_dict = {}
     settings = set_annotate_default_settings(settings={})
     
@@ -201,8 +253,8 @@ def generate_annotate_fields(frame):
 
     # Arrange input fields and labels
     for row, (name, data) in enumerate(vars_dict.items()):
-        ttk.Label(frame, text=f"{name.replace('_', ' ').capitalize()}:",
-                  background="black", foreground="white").grid(row=row, column=0)
+        tk.Label(frame, text=f"{name.replace('_', ' ').capitalize()}:", bg=style_out['bg_color'], fg=style_out['fg_color'], font=font_loader.get_font(size=font_size)).grid(row=row, column=0)
+        #ttk.Label(frame, text=f"{name.replace('_', ' ').capitalize()}:", background="black", foreground="white").grid(row=row, column=0)
         if isinstance(data['value'], list):
             # Convert lists to comma-separated strings
             data['entry'].insert(0, ','.join(map(str, data['value'])))
@@ -341,7 +393,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 +528,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 +542,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
@@ -125,7 +125,7 @@ def plot_image_mask_overlay(file, channels, cell_channel, nucleus_channel, patho
 
     return
 
-def plot_masks(batch, masks, flows, cmap='inferno', figuresize=20, nr=1, file_type='.npz', print_object_number=True):
+def plot_masks(batch, masks, flows, cmap='inferno', figuresize=10, nr=1, file_type='.npz', print_object_number=True):
     """
     Plot the masks and flows for a given batch of images.
 
@@ -476,7 +476,7 @@ def _filter_objects_in_plot(stack, cell_mask_dim, nucleus_mask_dim, pathogen_mas
 
     return stack
 
-def plot_arrays(src, figuresize=50, cmap='inferno', nr=1, normalize=True, q1=1, q2=99):
+def plot_arrays(src, figuresize=10, cmap='inferno', nr=1, normalize=True, q1=1, q2=99):
     """
     Plot randomly selected arrays from a given directory.
 
@@ -870,7 +870,7 @@ def _save_scimg_plot(src, nr_imgs=16, channel_indices=[0,1,2], um_per_pixel=0.1,
 
     return
 
-def _plot_cropped_arrays(stack, filename, figuresize=20, cmap='inferno', threshold=500):
+def _plot_cropped_arrays(stack, filename, figuresize=10, cmap='inferno', threshold=500):
     """
     Plot cropped arrays.
 
@@ -997,7 +997,7 @@ def _display_gif(path):
     with open(path, 'rb') as file:
         display(ipyimage(file.read()))
         
-def _plot_recruitment(df, df_type, channel_of_interest, target, columns=[], figuresize=50):
+def _plot_recruitment(df, df_type, channel_of_interest, target, columns=[], figuresize=10):
     """
     Plot recruitment data for different conditions and pathogens.
 
@@ -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()

spacr/sequencing.py

@@ -223,14 +223,10 @@ def save_to_hdf(queue, output_file, complevel=9, compression='zlib'):
     Save data from a queue to an HDF file.
 
     Parameters:
-    - queue: Queue object
-        The queue containing the data to be saved.
-    - output_file: strs
-        The path to the output HDF file.
-    - complevel: int, optional
-        The compression level to use (default is 9).
-    - compression: str, optional
-        The compression algorithm to use (default is 'zlib').
+    - queue: Queue object containing chunks of data to be saved
+    - output_file: Path to the output HDF file
+    - complevel: Compression level (default: 9)
+    - compression: Compression algorithm (default: 'zlib')
 
     Returns:
     None