dsa-helpers 1.1.18.dev0__tar.gz → 1.1.19.dev2__tar.gz

{dsa_helpers-1.1.18.dev0/dsa_helpers.egg-info → dsa_helpers-1.1.19.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dsa_helpers
-Version: 1.1.18.dev0
+Version: 1.1.19.dev2
 Summary: Utility functions for working with the DSA girder client.
 Author-email: "Juan C. Vizcarra" <jvizcar@emory.edu>
 Maintainer-email: "Juan C. Vizcarra" <jvizcar@emory.edu>

dsa_helpers-1.1.19.dev2/docker/Dockerfile

@@ -0,0 +1,12 @@
+FROM pytorch/pytorch:2.8.0-cuda12.6-cudnn9-runtime
+
+RUN pip install --no-cache-dir dsa-helpers && \
+    pip install --no-cache-dir large-image[all] --find-links https://girder.github.io/large_image_wheels && \
+    pip install --no-cache-dir histomicstk
+
+# # Dependencies for OpenCV (not needed anymore, but keeping for reference)
+# RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y
+
+RUN mkdir /app && chmod 777 /app && mkdir /data && chmod 777 /data
+
+WORKDIR /app

{dsa_helpers-1.1.18.dev0 → dsa_helpers-1.1.19.dev2}/dsa_helpers/__init__.py

@@ -3,7 +3,7 @@ from .imread import imread
 from .imwrite import imwrite
 
 # Version of the dsa-helpers package
-__version__ = "1.1.18dev0"
+__version__ = "1.1.19dev2"
 
 # To avoid slow downs, do not allow from dsa_helpers import * to import anything.
 __all__ = []

dsa_helpers-1.1.19.dev2/dsa_helpers/ml/callbacks.py

@@ -0,0 +1,134 @@
+"""Provides custom callbacks for the Hugging Face Trainer class."""
+
+from transformers import TrainerCallback
+from pathlib import Path
+import pandas as pd
+from copy import deepcopy
+import os, shutil
+
+
+class MetricsLoggerCallback(TrainerCallback):
+    def _load_current_file(self, output_dir):
+        output_dir = Path(output_dir)
+        metrics_file = output_dir / "metrics.csv"
+
+        return (
+            pd.read_csv(metrics_file)
+            if metrics_file.exists()
+            else pd.DataFrame()
+        )
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        """
+        Triggered during training logging (e.g., for training loss).
+        """
+        if "loss" in logs:  # Ensure training loss is available
+            metrics_df = self._load_current_file(args.output_dir)
+
+            # Append the epoch and the training loss.
+            loss = logs["loss"]
+            epoch = state.epoch
+
+            metrics_df = pd.concat(
+                [
+                    metrics_df,
+                    pd.DataFrame([{"epoch": epoch, "train_loss": loss}]),
+                ],
+                ignore_index=True,
+            )
+
+            # Save it back to file.
+            metrics_df.to_csv(
+                Path(args.output_dir) / "metrics.csv", index=False
+            )
+
+    def on_evaluate(self, args, state, control, metrics, **kwargs):
+        """
+        This is triggered at the end of evaluation (including at the end of each epoch).
+        """
+        # Get the current metrics file.
+        metrics_df = self._load_current_file(args.output_dir)
+
+        # Add metrics to the last row of the current metrics file.
+        for k, v in metrics.items():
+            metrics_df.at[metrics_df.index[-1], k] = v
+
+        # Save the updated metrics file.
+        metrics_df.to_csv(Path(args.output_dir) / "metrics.csv", index=False)
+
+
+class BestModelCallback(TrainerCallback):
+    """Callback for tracking and saving the best model checkpoint.
+
+    This class was developed with the help of ChatGPT."""
+
+    def __init__(self, save_dir: str):
+        """
+        Initialize the callback.
+
+        Args:
+            save_dir (str): Directory where model is saving checkpoints.
+                The best checkpoint will be saved in a subdirectory
+                called "best-checkpoint".
+
+        """
+        self.best_loss = float("inf")
+        self.best_model_state = None
+        self.save_dir = save_dir
+
+    def on_evaluate(self, args, state, control, metrics, model, **kwargs):
+        """Called after evaluation. Track the best model state."""
+        # Get the mean_dice metric
+        current_metric = metrics.get("eval_loss")
+
+        if current_metric is not None and current_metric < self.best_loss:
+            self.best_loss = current_metric
+            print("New best loss.")
+
+            # Store the model state.
+            self.best_model_state = deepcopy(model.state_dict())
+
+    def on_train_end(self, args, state, control, model, **kwargs):
+        """Logic:
+        * Create a best-checkpoint directory
+        * Save the best model state into the best-checkpoint directory
+        * Copy all the other model files (from the normal checkpoint
+          directory) into the best-checkpoint directory
+
+        """
+        if self.best_model_state is not None:
+            # Create best-checkpoint directory
+            best_checkpoint_dir = os.path.join(
+                self.save_dir, "best-checkpoint"
+            )
+            os.makedirs(best_checkpoint_dir, exist_ok=True)
+
+            # Load the best model state
+            model.load_state_dict(self.best_model_state)
+
+            # Save the model (this creates model.safetensors or pytorch_model.bin)
+            model.save_pretrained(best_checkpoint_dir)
+
+            # Copy other checkpoint files from the latest checkpoint
+            # Find the latest checkpoint directory (checkpoint-{step_number})
+            save_path = Path(self.save_dir)
+            checkpoint_dirs = [
+                d
+                for d in save_path.iterdir()
+                if d.is_dir() and d.name.startswith("checkpoint-")
+            ]
+            if checkpoint_dirs:
+                # Sort by step number to get the latest
+                latest_checkpoint = max(
+                    checkpoint_dirs, key=lambda x: int(x.name.split("-")[1])
+                )
+
+                # Copy all files from the latest checkpoint, except the model files we just created
+                for file_path in latest_checkpoint.iterdir():
+                    if file_path.is_file():
+                        # Skip model files since we already saved the best model state
+                        if not file_path.name.startswith(
+                            ("model.", "pytorch_model.")
+                        ):
+                            dst = Path(best_checkpoint_dir) / file_path.name
+                            shutil.copy2(file_path, dst)

{dsa_helpers-1.1.18.dev0 → dsa_helpers-1.1.19.dev2}/dsa_helpers/ml/segformer_semantic_segmentation/train.py

@@ -6,6 +6,7 @@ from transformers import (
 )
 import numpy as np
 from pathlib import Path
+import tempfile, shutil
 
 import torch
 from torch import nn
@@ -13,7 +14,7 @@ from torch import nn
 from .utils import create_segformer_segmentation_dataset
 from .transforms import get_train_transforms, val_transforms
 from .evaluate import per_class_dice_on_dataset
-from ..callbacks import MetricsLoggerCallback
+from ..callbacks import MetricsLoggerCallback, BestModelCallback
 
 
 def train(
@@ -41,6 +42,7 @@ def train(
     contrast: float = 0.25,
     saturation: float = 0.25,
     hue: float = 0.1,
+    rank: int = 0,
 ):
     """Workflow for training a SegFormer semantic segmentation model.
 
@@ -101,6 +103,7 @@ def train(
             images. Defaults to 0.25.
         hue (float, optional): Hue factor to apply to the images.
             Defaults to 0.1.
+        rank (int, optional): Rank of the process. Defaults to 0.
 
     Returns:
         tuple: A tuple containing the trainer and the results dictionary.
@@ -108,8 +111,11 @@ def train(
     """
     save_dir_path = Path(save_dir)
 
-    if save_dir_path.is_dir():
-        raise FileExistsError(f"Save directory {save_dir} already exists.")
+    if rank == 0:
+        if save_dir_path.is_dir():
+            raise FileExistsError(f"Save directory {save_dir} already exists.")
+        else:
+            save_dir_path.mkdir(parents=True, exist_ok=True)
 
     # Read the dataframes if they are strings.
     if isinstance(train_data, str):
@@ -139,6 +145,10 @@ def train(
             ignore_mismatched_sizes=True,
         ).to(device)
 
+    if rank != 0:
+        # Create temp directory for the model.
+        save_dir = tempfile.mkdtemp()
+
     # Training arguments.
     training_args = TrainingArguments(
         str(save_dir),
@@ -150,7 +160,6 @@ def train(
         eval_strategy="epoch",
         save_strategy="epoch",
         logging_strategy="epoch",
-        load_best_model_at_end=True,
         save_total_limit=1,
     )
 
@@ -236,6 +245,11 @@ def train(
         metrics["mean_dice"] = float(np.mean(list(metrics.values())))
         return metrics
 
+    if rank == 0:
+        callbacks = [MetricsLoggerCallback, BestModelCallback(str(save_dir))]
+    else:
+        callbacks = [MetricsLoggerCallback]
+
     # Trainer
     trainer = Trainer(
         model=model,
@@ -243,73 +257,62 @@ def train(
         train_dataset=train_dataset,
         eval_dataset=val_dataset,
         compute_metrics=compute_metrics,
-        callbacks=[MetricsLoggerCallback],  # Add the custom callback here
+        callbacks=callbacks,
     )
 
-    # Train the model.
-    _ = trainer.train()
-
-    # Rename the checkpoint to not include a step number.
-    save_dir = Path(save_dir)
-    for p in save_dir.iterdir():
-        if p.is_dir() and p.name.startswith("checkpoint-"):
-            p.rename(save_dir / "checkpoint")
+    if rank != 0:
+        # Train the model.
+        _ = trainer.train()
 
-    results = {"train": {}, "val": {}}
-
-    for log in trainer.state.log_history[:-1]:
-        epoch = log["epoch"]
+        # Delete the temp directory.
+        shutil.rmtree(save_dir)
+        return None
+    else:
+        print("Starting training...")
+        _ = trainer.train()
+        print("Training completed successfully!")
 
-        if epoch % 1:
-            continue  # only track end of each epoch
+        results = {"train": {}, "val": {}}
 
-        # Check if this is tracking train or validation.
-        if "eval_loss" in log:
-            data = results["val"]
-        else:
-            data = results["train"]
+        for log in trainer.state.log_history[:-1]:
+            epoch = log["epoch"]
 
-        for k, v in log.items():
-            if k.startswith("eval_"):
-                k = k[5:]
+            if epoch % 1:
+                continue  # only track end of each epoch
 
-            if k not in data:
-                data[k] = []
+            # Check if this is tracking train or validation.
+            if "eval_loss" in log:
+                data = results["val"]
+            else:
+                data = results["train"]
 
-            data[k].append(v)
+            for k, v in log.items():
+                if k.startswith("eval_"):
+                    k = k[5:]
 
-    # For each of the datasets provided, get the metrics.
-    model = trainer.model
+                if k not in data:
+                    data[k] = []
 
-    print("Calculating train metrics...")
-    results["train"]["metrics"] = per_class_dice_on_dataset(
-        model,
-        train_data,
-        label2id,
-        batch_size=batch_size,
-        device=device,
-        tile_size=tile_size,
-        tqdm_notebook=tqdm_notebook,
-    )
+                data[k].append(v)
 
-    print("Calculating val metrics...")
-    results["val"]["metrics"] = per_class_dice_on_dataset(
-        model,
-        val_data,
-        label2id,
-        batch_size=batch_size,
-        device=device,
-        tile_size=tile_size,
-        tqdm_notebook=tqdm_notebook,
-    )
+        # For each of the datasets provided, get the metrics.
+        model = trainer.model
 
-    if test_data is not None:
-        print("Calculating test metrics...")
-        results["test"] = {}
+        print("Calculating train metrics...")
+        results["train"]["metrics"] = per_class_dice_on_dataset(
+            model,
+            train_data,
+            label2id,
+            batch_size=batch_size,
+            device=device,
+            tile_size=tile_size,
+            tqdm_notebook=tqdm_notebook,
+        )
 
-        results["test"]["metrics"] = per_class_dice_on_dataset(
+        print("Calculating val metrics...")
+        results["val"]["metrics"] = per_class_dice_on_dataset(
             model,
-            test_data,
+            val_data,
             label2id,
             batch_size=batch_size,
             device=device,
@@ -317,4 +320,18 @@ def train(
             tqdm_notebook=tqdm_notebook,
         )
 
-    return trainer, results
+        if test_data is not None:
+            print("Calculating test metrics...")
+            results["test"] = {}
+
+            results["test"]["metrics"] = per_class_dice_on_dataset(
+                model,
+                test_data,
+                label2id,
+                batch_size=batch_size,
+                device=device,
+                tile_size=tile_size,
+                tqdm_notebook=tqdm_notebook,
+            )
+
+        return trainer, results

{dsa_helpers-1.1.18.dev0 → dsa_helpers-1.1.19.dev2/dsa_helpers.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dsa_helpers
-Version: 1.1.18.dev0
+Version: 1.1.19.dev2
 Summary: Utility functions for working with the DSA girder client.
 Author-email: "Juan C. Vizcarra" <jvizcar@emory.edu>
 Maintainer-email: "Juan C. Vizcarra" <jvizcar@emory.edu>

{dsa_helpers-1.1.18.dev0 → dsa_helpers-1.1.19.dev2}/pyproject.toml

@@ -41,7 +41,7 @@ classifiers = [
     "License :: OSI Approved :: MIT License",
 ]
 keywords = ["dsa", "girder", "girder_client"]
-version = "v1.1.18dev0"
+version = "v1.1.19dev2"
 
 [project.optional-dependencies]
 opencv = ["opencv-python"]

dsa_helpers-1.1.18.dev0/docker/Dockerfile

@@ -1,14 +0,0 @@
-FROM pytorch/pytorch:2.8.0-cuda12.9-cudnn9-runtime
-
-RUN pip install --no-cache-dir dsa-helpers
-RUN pip install --no-cache-dir large-image[all] --find-links https://girder.github.io/large_image_wheels
-RUN python -m pip install --no-cache-dir histomicstk --find-links https://girder.github.io/large_image_wheels
-
-# Dependencies for OpenCV.
-RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y
-
-RUN mkdir /app
-RUN chmod 777 /app
-
-RUN mkdir /data
-RUN chmod 777 /data

dsa_helpers-1.1.18.dev0/dsa_helpers/ml/callbacks.py

@@ -1,55 +0,0 @@
-"""Provides custom callbacks for the Hugging Face Trainer class."""
-
-from transformers import TrainerCallback
-from pathlib import Path
-import pandas as pd
-
-
-class MetricsLoggerCallback(TrainerCallback):
-    def _load_current_file(self, output_dir):
-        output_dir = Path(output_dir)
-        metrics_file = output_dir / "metrics.csv"
-
-        return (
-            pd.read_csv(metrics_file)
-            if metrics_file.exists()
-            else pd.DataFrame()
-        )
-
-    def on_log(self, args, state, control, logs=None, **kwargs):
-        """
-        Triggered during training logging (e.g., for training loss).
-        """
-        if "loss" in logs:  # Ensure training loss is available
-            metrics_df = self._load_current_file(args.output_dir)
-
-            # Append the epoch and the training loss.
-            loss = logs["loss"]
-            epoch = state.epoch
-
-            metrics_df = pd.concat(
-                [
-                    metrics_df,
-                    pd.DataFrame([{"epoch": epoch, "train_loss": loss}]),
-                ],
-                ignore_index=True,
-            )
-
-            # Save it back to file.
-            metrics_df.to_csv(
-                Path(args.output_dir) / "metrics.csv", index=False
-            )
-
-    def on_evaluate(self, args, state, control, metrics, **kwargs):
-        """
-        This is triggered at the end of evaluation (including at the end of each epoch).
-        """
-        # Get the current metrics file.
-        metrics_df = self._load_current_file(args.output_dir)
-
-        # Add metrics to the last row of the current metrics file.
-        for k, v in metrics.items():
-            metrics_df.at[metrics_df.index[-1], k] = v
-
-        # Save the updated metrics file.
-        metrics_df.to_csv(Path(args.output_dir) / "metrics.csv", index=False)