celldetective 1.4.2__py3-none-any.whl → 1.5.0b0__py3-none-any.whl

celldetective/processes/train_signal_model.py

@@ -0,0 +1,265 @@
+from multiprocessing import Process
+import time
+import os
+import json
+from glob import glob
+import numpy as np
+from art import tprint
+from tensorflow.python.keras.callbacks import Callback
+
+from celldetective.signals import SignalDetectionModel
+from celldetective.log_manager import get_logger
+from celldetective.utils.model_loaders import locate_signal_model
+
+logger = get_logger(__name__)
+
+
+class ProgressCallback(Callback):
+
+    def __init__(self, queue=None, total_epochs=100, stop_event=None):
+        super().__init__()
+        self.queue = queue
+        self.total_epochs = total_epochs
+        self.current_step = 0
+        self.t0 = time.time()
+        self.stop_event = stop_event
+
+    def on_epoch_begin(self, epoch, logs=None):
+        self.epoch_start_time = time.time()
+
+    def on_batch_end(self, batch, logs=None):
+        # Update frame bar (bottom bar) for batch progress
+        # logs has 'size' and 'batch'
+        # We need total batches. Keras doesn't always pass it easily in logs unless we know steps_per_epoch.
+        # But self.params['steps'] should have it if available.
+        if self.params and "steps" in self.params:
+            total_steps = self.params["steps"]
+            batch_progress = ((batch + 1) / total_steps) * 100
+            if self.queue is not None:
+                # Send generic batch update (frequent)
+                self.queue.put(
+                    {
+                        "frame_progress": batch_progress,
+                        "frame_time": f"Batch {batch + 1}/{total_steps}",
+                    }
+                )
+
+    def on_epoch_end(self, epoch, logs=None):
+        if self.stop_event and self.stop_event.is_set():
+            logger.info("Interrupting training...")
+            self.model.stop_training = True
+            self.stop_event.clear()
+
+        self.current_step += 1
+        # Send signal for progress bar
+        sum_done = (self.current_step) / self.total_epochs * 100
+        mean_exec_per_step = (time.time() - self.t0) / (self.current_step)
+        pred_time = (self.total_epochs - self.current_step) * mean_exec_per_step
+
+        # Format time string
+        if pred_time > 60:
+            time_str = f"{pred_time/60:.1f} min"
+        else:
+            time_str = f"{pred_time:.1f} s"
+
+        if self.queue is not None:
+            # Update Position bar (middle) for Epoch progress
+            msg = {
+                "pos_progress": sum_done,
+                "pos_time": f"Epoch {self.current_step}/{self.total_epochs} (ETA: {time_str})",
+                "frame_progress": 0,  # Reset batch bar
+                "frame_time": "Batch 0/0",
+            }
+            # Attempt to extract metrics for plotting
+            if logs:
+                # Infer model type
+                if "iou" in logs:
+                    model_name = "Classifier"
+                else:
+                    model_name = "Regressor"
+
+                # Send all scalar metrics
+                msg["plot_data"] = {
+                    "epoch": epoch + 1,  # 1-based for plot
+                    "metrics": {
+                        k: float(v) for k, v in logs.items() if not k.startswith("val_")
+                    },
+                    "val_metrics": {
+                        k: float(v) for k, v in logs.items() if k.startswith("val_")
+                    },
+                    "model_name": model_name,
+                    "total_epochs": self.params.get("epochs", self.total_epochs),
+                }
+            self.queue.put(msg)
+
+    def on_training_result(self, result):
+        if self.queue is not None:
+            self.queue.put({"training_result": result})
+
+
+class TrainSignalModelProcess(Process):
+
+    def __init__(self, queue=None, process_args=None, *args, **kwargs):
+
+        super().__init__(*args, **kwargs)
+
+        self.queue = queue
+
+        if process_args is not None:
+            for key, value in process_args.items():
+                setattr(self, key, value)
+
+        tprint("Train event detection")
+        self.read_instructions()
+        self.extract_training_params()
+
+        self.sum_done = 0
+        self.t0 = time.time()
+
+    def read_instructions(self):
+
+        if os.path.exists(self.instructions):
+            with open(self.instructions, "r") as f:
+                self.training_instructions = json.load(f)
+        else:
+            logger.error("Training instructions could not be found. Abort.")
+            self.abort_process()
+
+        all_classes = []
+        for d in self.training_instructions["ds"]:
+            datasets = glob(d + os.sep + "*.npy")
+            for dd in datasets:
+                data = np.load(dd, allow_pickle=True)
+                classes = np.unique([ddd["class"] for ddd in data])
+                all_classes.extend(classes)
+        all_classes = np.unique(all_classes)
+        n_classes = len(all_classes)
+
+        self.model_params = {
+            k: self.training_instructions[k]
+            for k in (
+                "pretrained",
+                "model_signal_length",
+                "channel_option",
+                "n_channels",
+                "label",
+            )
+            if k in self.training_instructions
+        }
+        self.model_params.update({"n_classes": n_classes})
+        self.train_params = {
+            k: self.training_instructions[k]
+            for k in (
+                "model_name",
+                "target_directory",
+                "channel_option",
+                "recompile_pretrained",
+                "test_split",
+                "augment",
+                "epochs",
+                "learning_rate",
+                "batch_size",
+                "validation_split",
+                "normalization_percentile",
+                "normalization_values",
+                "normalization_clip",
+            )
+            if k in self.training_instructions
+        }
+
+    def neighborhood_postprocessing(self):
+
+        # if neighborhood of interest in training instructions, write it in config!
+        if "neighborhood_of_interest" in self.training_instructions:
+            if self.training_instructions["neighborhood_of_interest"] is not None:
+
+                model_path = locate_signal_model(
+                    self.training_instructions["model_name"], path=None, pairs=True
+                )
+                complete_path = model_path  # +model
+                complete_path = rf"{complete_path}"
+                model_config_path = os.sep.join([complete_path, "config_input.json"])
+                model_config_path = rf"{model_config_path}"
+
+                f = open(model_config_path)
+                config = json.load(f)
+                config.update(
+                    {
+                        "neighborhood_of_interest": self.training_instructions[
+                            "neighborhood_of_interest"
+                        ],
+                        "reference_population": self.training_instructions[
+                            "reference_population"
+                        ],
+                        "neighbor_population": self.training_instructions[
+                            "neighbor_population"
+                        ],
+                    }
+                )
+                json_string = json.dumps(config)
+                with open(model_config_path, "w") as outfile:
+                    outfile.write(json_string)
+
+    def run(self):
+        self.queue.put({"status": "Loading datasets..."})
+        model = SignalDetectionModel(**self.model_params)
+
+        total_epochs = self.train_params["epochs"] * 3
+        cb = ProgressCallback(
+            queue=self.queue,
+            total_epochs=total_epochs,
+            stop_event=getattr(self, "stop_event", None),
+        )
+
+        model.fit_from_directory(
+            self.training_instructions["ds"], callbacks=[cb], **self.train_params
+        )
+
+        # Send results to GUI
+        if hasattr(model, "dico"):
+            result_keys = [
+                "val_confusion",
+                "test_confusion",
+                "val_predictions",
+                "val_ground_truth",
+                "test_predictions",
+                "test_ground_truth",
+                "val_mse",
+            ]
+            results = {k: model.dico[k] for k in result_keys if k in model.dico}
+            # Only send if we have something relevant
+            if results:
+                self.queue.put({"training_result": results})
+
+        self.neighborhood_postprocessing()
+        self.queue.put("finished")
+        self.queue.close()
+
+    def extract_training_params(self):
+
+        self.training_instructions.update(
+            {"n_channels": len(self.training_instructions["channel_option"])}
+        )
+        self.model_params["n_channels"] = self.training_instructions["n_channels"]
+        if self.training_instructions["augmentation_factor"] > 1.0:
+            self.training_instructions.update({"augment": True})
+        else:
+            self.training_instructions.update({"augment": False})
+        self.training_instructions.update({"test_split": 0.0})
+
+    def end_process(self):
+
+        # self.terminate()
+
+        # if self.model_type=="stardist_utils":
+        # 	from stardist_utils.models import StarDist2D
+        # 	self.model = StarDist2D(None, name=self.model_name, basedir=self.target_directory)
+        # 	self.model.optimize_thresholds(self.X_val,self.Y_val)
+
+        self.terminate()
+        self.queue.put("finished")
+
+    def abort_process(self):
+
+        self.terminate()
+        self.queue.put("error")

celldetective/processes/unified_process.py

@@ -0,0 +1,292 @@
+import time
+import os
+import gc
+from multiprocessing import Process
+from pathlib import Path
+
+from celldetective.log_manager import get_logger
+
+logger = get_logger(__name__)
+
+
+class UnifiedBatchProcess(Process):
+    """
+    A unified process that handles Segmentation, Tracking, Measurement, and Signal Analysis
+    in a sequential manner for each position, updating a multi-bar progress window.
+    """
+
+    def __init__(self, queue, process_args=None):
+        super(UnifiedBatchProcess, self).__init__()
+        self.queue = queue
+        self.process_args = process_args
+        self.batch_structure = process_args.get("batch_structure", {})
+
+        self.run_segmentation = process_args.get("run_segmentation", False)
+        self.run_tracking = process_args.get("run_tracking", False)
+        self.run_measurement = process_args.get("run_measurement", False)
+        self.run_signals = process_args.get("run_signals", False)
+
+        self.seg_args = process_args.get("seg_args", {})
+        self.track_args = process_args.get("track_args", {})
+        self.measure_args = process_args.get("measure_args", {})
+        self.signal_args = process_args.get("signal_args", {})
+        self.log_file = process_args.get("log_file", None)
+
+    def run(self):
+
+        if self.log_file is not None:
+            from celldetective.log_manager import setup_logging
+
+            setup_logging(self.log_file)
+
+        logger.info("Starting Unified Batch Process...")
+
+        # Initialize Workers
+        # Propagate batch structure to sub-processes so they can locate experiment config
+        for args_dict in [
+            self.seg_args,
+            self.track_args,
+            self.measure_args,
+            self.signal_args,
+        ]:
+            if isinstance(args_dict, dict):
+                args_dict["batch_structure"] = self.batch_structure
+
+        seg_worker = None
+        model = None
+        scale_model = None
+
+        if self.run_segmentation:
+            logger.info("Initializing the segmentation worker...")
+            self.queue.put({"status": "Initializing segmentation..."})
+
+            if "threshold_instructions" in self.seg_args:
+                from celldetective.processes.segment_cells import (
+                    SegmentCellThresholdProcess,
+                )
+
+                seg_worker = SegmentCellThresholdProcess(
+                    queue=self.queue, process_args=self.seg_args
+                )
+            else:
+                from celldetective.processes.segment_cells import SegmentCellDLProcess
+
+                seg_worker = SegmentCellDLProcess(
+                    queue=self.queue, process_args=self.seg_args
+                )
+
+                if seg_worker.model_type == "stardist":
+                    logger.info("Loading the StarDist library...")
+                    from celldetective.utils.stardist_utils import _prep_stardist_model
+
+                    model, scale_model = _prep_stardist_model(
+                        seg_worker.model_name,
+                        Path(seg_worker.model_complete_path).parent,
+                        use_gpu=seg_worker.use_gpu,
+                        scale=seg_worker.scale,
+                    )
+                elif seg_worker.model_type == "cellpose":
+                    logger.info("Loading the cellpose_utils library...")
+                    from celldetective.utils.cellpose_utils import _prep_cellpose_model
+
+                    model, scale_model = _prep_cellpose_model(
+                        seg_worker.model_name,
+                        seg_worker.model_complete_path,
+                        use_gpu=seg_worker.use_gpu,
+                        n_channels=len(seg_worker.required_channels),
+                        scale=seg_worker.scale,
+                    )
+
+        track_worker = None
+        if self.run_tracking:
+            from celldetective.processes.track_cells import TrackingProcess
+
+            logger.info("Initializing the tracking worker...")
+            self.queue.put({"status": "Initializing tracking..."})
+            track_worker = TrackingProcess(
+                queue=self.queue, process_args=self.track_args
+            )
+
+        measure_worker = None
+        if self.run_measurement:
+            logger.info("Loading the measurement libraries...")
+            from celldetective.processes.measure_cells import MeasurementProcess
+
+            logger.info("Initializing the measurement worker...")
+            self.queue.put({"status": "Initializing measurements..."})
+            measure_worker = MeasurementProcess(
+                queue=self.queue, process_args=self.measure_args
+            )
+
+        signal_worker = None
+        signal_model = None
+
+        if self.run_signals:
+            from celldetective.utils.event_detection import _prep_event_detection_model
+
+            try:
+                logger.info("Loading the event detection model...")
+                self.queue.put({"status": "Loading event detection model..."})
+                model_name = self.signal_args["model_name"]
+                signal_model = _prep_event_detection_model(
+                    model_name, use_gpu=self.signal_args.get("gpu", True)
+                )
+            except Exception as e:
+                logger.error(
+                    f"Failed to initialize event detection model: {e}", exc_info=True
+                )
+                self.run_signals = False  # Disable signal analysis if model fails
+
+        if self.run_signals:
+            from celldetective.processes.detect_events import SignalAnalysisProcess
+
+            logger.info("Initializing the event detection worker...")
+            signal_worker = SignalAnalysisProcess(
+                queue=self.queue, process_args=self.signal_args
+            )
+            signal_worker.signal_model_instance = signal_model
+
+        self.t0_well = time.time()
+
+        for w_i, (w_idx, well_data) in enumerate(self.batch_structure.items()):
+
+            positions = well_data["positions"]
+
+            # Well Progress Update
+            elapsed = time.time() - self.t0_well
+            if w_i > 0:
+                avg = elapsed / w_i
+                rem = (len(self.batch_structure) - w_i) * avg
+                mins = int(rem // 60)
+                secs = int(rem % 60)
+                well_str = f"Well {w_i + 1}/{len(self.batch_structure)} - {mins} m {secs} s left"
+            else:
+                well_str = f"Processing well {w_i + 1}/{len(self.batch_structure)}..."
+
+            self.queue.put(
+                {
+                    "well_progress": (w_i / len(self.batch_structure)) * 100,
+                    "well_time": well_str,
+                }
+            )
+
+            self.t0_pos = time.time()
+
+            for pos_idx, pos_path in enumerate(positions):
+
+                # Position Progress Update
+                elapsed_pos = time.time() - self.t0_pos
+                if pos_idx > 0:
+                    avg_p = elapsed_pos / pos_idx
+                    rem_p = (len(positions) - pos_idx) * avg_p
+                    mins_p = int(rem_p // 60)
+                    secs_p = int(rem_p % 60)
+                    pos_str = f"Pos {pos_idx + 1}/{len(positions)} - {mins_p} m {secs_p} s left"
+                else:
+                    pos_str = f"Processing position {pos_idx + 1}/{len(positions)}..."
+
+                self.queue.put(
+                    {
+                        "pos_progress": (pos_idx / len(positions)) * 100,
+                        "pos_time": pos_str,
+                    }
+                )
+
+                # Calculate active steps for this run
+                active_steps = []
+                if self.run_segmentation:
+                    active_steps.append("Segmentation")
+                if self.run_tracking:
+                    active_steps.append("Tracking")
+                if self.run_measurement:
+                    active_steps.append("Measurement")
+                if self.run_signals:
+                    active_steps.append("Event detection")
+
+                total_steps = len(active_steps)
+                current_step = 0
+
+                try:
+                    # --- SEGMENTATION ---
+                    if self.run_segmentation and seg_worker:
+                        current_step += 1
+                        step_info = f"[Step {current_step}/{total_steps}]"
+                        msg = f"{step_info} Segmenting {os.path.basename(pos_path)}..."
+                        logger.info(msg)
+                        self.queue.put({"status": msg})
+
+                        seg_worker.setup_for_position(pos_path)
+
+                        if not "threshold_instructions" in self.seg_args:
+                            seg_worker.process_position(
+                                model=model, scale_model=scale_model
+                            )
+                        else:
+                            seg_worker.process_position()
+
+                    # --- TRACKING ---
+                    if self.run_tracking and track_worker:
+                        current_step += 1
+                        step_info = f"[Step {current_step}/{total_steps}]"
+                        msg = f"{step_info} Tracking {os.path.basename(pos_path)}..."
+                        logger.info(msg)
+                        self.queue.put({"status": msg})
+
+                        track_worker.setup_for_position(pos_path)
+                        track_worker.process_position()
+
+                    # --- MEASUREMENT ---
+                    if self.run_measurement and measure_worker:
+                        current_step += 1
+                        step_info = f"[Step {current_step}/{total_steps}]"
+                        msg = f"{step_info} Measuring {os.path.basename(pos_path)}..."
+                        logger.info(msg)
+                        self.queue.put({"status": msg})
+
+                        measure_worker.setup_for_position(pos_path)
+                        measure_worker.process_position()
+
+                    # --- SIGNAL ANALYSIS ---
+                    if self.run_signals and signal_worker:
+                        current_step += 1
+                        step_info = f"[Step {current_step}/{total_steps}]"
+                        msg = f"{step_info} Detecting events in position {os.path.basename(pos_path)}..."
+                        logger.info(msg)
+                        self.queue.put({"status": msg})
+
+                        signal_worker.setup_for_position(pos_path)
+                        signal_worker.process_position(model=signal_model)
+
+                except Exception as e:
+                    logger.error(f"Error processing position {pos_path}: {e}")
+                    self.queue.put(
+                        {
+                            "status": f"Error at {os.path.basename(pos_path)}. Skipping..."
+                        }
+                    )
+                    logger.error(
+                        f"Skipping position {os.path.basename(pos_path)} due to error: {e}",
+                        exc_info=True,
+                    )
+                    continue
+
+                gc.collect()
+
+                # Update Position Progress (Complete)
+                self.queue.put({"pos_progress": ((pos_idx + 1) / len(positions)) * 100})
+
+            # Update Well Progress (Complete)
+            self.queue.put(
+                {"well_progress": ((w_i + 1) / len(self.batch_structure)) * 100}
+            )
+
+        self.queue.put("finished")
+        self.queue.close()
+
+    def end_process(self):
+        try:
+            if self.is_alive():
+                self.terminate()
+                self.join()
+        except Exception as e:
+            logger.error(f"Error terminating process: {e}")