Merge pull request #18 from MilagrosMarin/update_docs

feat: Add support to generate PNG version of fitting progress plots & refactor

Author: ttngu207

CHANGELOG.md

@@ -3,6 +3,10 @@
 Observes [Semantic Versioning](https://semver.org/spec/v2.0.0.html) standard and
 [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) convention.
 
+## [1.0.1] - 2025-09-23
++ Feat - Add support to generate PNG version of fitting progress plots in `PreFit`, `FullFit`, and `moseq_report` schema
++ Fix - Update path handling to use `Path` objects and `dj.logger`
+
 ## [1.0.0] - 2025-09-10
 
 > **BREAKING CHANGES** - This version contains breaking changes due to keypoint-moseq upgrade and API refactoring. Please review the changes below and update your code accordingly.

element_moseq/moseq_report.py

@@ -198,43 +198,67 @@ class PreFitReport(dj.Imported):
     -> moseq_train.PreFit
     ---
     fitting_progress_pdf: attach # fitting_progress.pdf
+    fitting_progress_png: attach # fitting_progress.png
     """
 
     def make(self, key):
         prefit_model_name = (moseq_train.PreFit & key).fetch1("model_name")
         prefit_model_dir = find_full_path(
             moseq_train.get_kpms_processed_data_dir(), prefit_model_name
         )
-        prefit_output_dir = Path(prefit_model_dir) / "fitting_progress.pdf"
-        if prefit_output_dir.exists():
-            self.insert1({**key, "fitting_progress_pdf": prefit_output_dir})
-        else:
+
+        pdf_path = prefit_model_dir / "fitting_progress.pdf"
+        png_path = prefit_model_dir / "fitting_progress.png"
+
+        if not pdf_path.exists():
+            raise FileNotFoundError(
+                f"PreFit PDF progress plot not found at {pdf_path}. "
+            )
+
+        if not png_path.exists():
             raise FileNotFoundError(
-                f"PreFit fitting_progress.pdf not found at {prefit_output_dir}"
+                f"PreFit PNG progress plot not found at {png_path}. "
             )
 
+        # Both files exist, insert them
+        self.insert1(
+            {**key, "fitting_progress_pdf": pdf_path, "fitting_progress_png": png_path}
+        )
+
 
 @schema
 class FullFitReport(dj.Imported):
     definition = """
     -> moseq_train.FullFit
     ---
     fitting_progress_pdf: attach # fitting_progress.pdf
+    fitting_progress_png: attach # fitting_progress.png
     """
 
     def make(self, key):
         fullfit_model_name = (moseq_train.FullFit & key).fetch1("model_name")
         fullfit_model_dir = find_full_path(
             moseq_train.get_kpms_processed_data_dir(), fullfit_model_name
         )
-        fullfit_output_file = Path(fullfit_model_dir) / "fitting_progress.pdf"
-        if fullfit_output_file.exists():
-            self.insert1({**key, "fitting_progress_pdf": fullfit_output_file})
-        else:
+
+        pdf_path = fullfit_model_dir / "fitting_progress.pdf"
+        png_path = fullfit_model_dir / "fitting_progress.png"
+
+        if not pdf_path.exists():
+            raise FileNotFoundError(
+                f"FullFit PDF progress plot not found at {pdf_path}. "
+            )
+
+        if not png_path.exists():
             raise FileNotFoundError(
-                f"FullFit fitting_progress.pdf not found at {fullfit_output_file}"
+                f"FullFit PNG progress plot not found at {png_path}. "
             )
 
+        # Both files exist, insert them
+        self.insert1(
+            {**key, "fitting_progress_pdf": pdf_path, "fitting_progress_png": png_path}
+        )
+
 
 @schema
 class InferenceReport(dj.Imported):

element_moseq/moseq_train.py

@@ -5,7 +5,6 @@
 
 import importlib
 import inspect
-import os
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Optional
@@ -15,6 +14,7 @@
 import numpy as np
 from element_interface.utils import find_full_path
 
+from .plotting import viz_utils
 from .readers import kpms_reader
 
 schema = dj.schema()
@@ -790,6 +790,8 @@ def make(self, key):
             end_time = datetime.now(timezone.utc)
 
             duration_seconds = (end_time - start_time).total_seconds()
+            viz_utils.copy_pdf_to_png(kpms_project_output_dir, model_name)
+
         else:
             duration_seconds = None
 
@@ -942,6 +944,8 @@ def make(self, key):
                 project_dir=kpms_project_output_dir.as_posix(),
                 ar_only=False,
                 num_iters=full_num_iterations,
+                generate_progress_plots=True,  # saved to {project_dir}/{model_name}/plots/
+                save_every_n_iters=25,
             )
             end_time = datetime.now(timezone.utc)
             duration_seconds = (end_time - start_time).total_seconds()
@@ -950,6 +954,7 @@ def make(self, key):
                 project_dir=kpms_project_output_dir.as_posix(),
                 model_name=Path(model_name).parts[-1],
             )
+            viz_utils.copy_pdf_to_png(kpms_project_output_dir, model_name)
 
         else:
             duration_seconds = None

element_moseq/plotting/viz_utils.py

@@ -7,12 +7,15 @@
 from textwrap import fill
 from typing import Dict, List, Optional, Tuple
 
+import datajoint as dj
 import matplotlib.pyplot as plt
 import numpy as np
 from jax_moseq.models.keypoint_slds import center_embedding
 from keypoint_moseq.util import get_distance_to_medoid, get_edges, plot_keypoint_traces
 from keypoint_moseq.viz import plot_pcs_3D
 
+logger = dj.logger
+
 _DLC_SUFFIX_RE = re.compile(
     r"(?:DLC_[A-Za-z0-9]+[A-Za-z]+(?:\d+)?(?:[A-Za-z]+)?"  # scorer-ish token
     r"(?:\w+)*)"  # optional extra blobs
@@ -158,7 +161,9 @@ def plot_medoid_distance_outliers(
     fig.savefig(plot_path, dpi=300)
 
     plt.close()
-    print(f"Saved keypoint distance outlier plot for {recording_name} to {plot_path}.")
+    logger.info(
+        f"Saved keypoint distance outlier plot for {recording_name} to {plot_path}."
+    )
     return fig
 
 
@@ -328,3 +333,41 @@ def plot_pcs(
             line_width * 2,
         )
     return fig
+
+
+def copy_pdf_to_png(project_dir, model_name):
+    """
+    Convert PDF progress plot to PNG format using pdf2image.
+    The fit_model function generates a single fitting_progress.pdf file.
+    This function should always succeed if the PDF exists.
+
+    Args:
+        project_dir (Path): Project directory path
+        model_name (str): Model name directory
+
+    Returns:
+        bool: True if conversion was successful, False otherwise
+
+    Raises:
+        FileNotFoundError: If the PDF file doesn't exist
+        RuntimeError: If conversion fails
+    """
+    from pdf2image import convert_from_path
+
+    # Construct paths for PDF and PNG files
+    model_dir = Path(project_dir) / model_name
+    pdf_path = model_dir / "fitting_progress.pdf"
+    png_path = model_dir / "fitting_progress.png"
+
+    # Check if PDF exists
+    if not pdf_path.exists():
+        raise FileNotFoundError(f"PDF progress plot not found at {pdf_path}")
+
+    # Convert PDF to PNG
+    images = convert_from_path(pdf_path, dpi=300)
+    if not images:
+        raise ValueError(f"No PDF file found at {pdf_path}")
+
+    images[0].save(png_path, "PNG")
+    logger.info(f"Generated PNG progress plot at {png_path}")
+    return True

element_moseq/readers/kpms_reader.py

@@ -1,13 +1,12 @@
-import logging
 import os
 from pathlib import Path
 from typing import Any, Dict, Union
 
+import datajoint as dj
 import jax.numpy as jnp
 import yaml
 
-logger = logging.getLogger("datajoint")
-
+logger = dj.logger
 
 DJ_CONFIG = "kpms_dj_config.yml"
 BASE_CONFIG = "config.yml"
@@ -72,11 +71,11 @@ def dj_generate_config(project_dir: str, **kwargs) -> str:
     base_cfg_path = _base_config_path(project_dir)
     dj_cfg_path = _dj_config_path(project_dir)
 
-    if os.path.exists(dj_cfg_path):
+    if Path(dj_cfg_path).exists():
         with open(dj_cfg_path, "r") as f:
             cfg = yaml.safe_load(f) or {}
     else:
-        if not os.path.exists(base_cfg_path):
+        if not Path(base_cfg_path).exists():
             raise FileNotFoundError(
                 f"Missing base config at {base_cfg_path}. Run upstream setup_project first. "
                 f"Expected either config.yml or config.yaml in {project_dir}."
@@ -105,7 +104,7 @@ def load_kpms_dj_config(
         indexing into 'use_bodyparts' by order.
     """
     dj_cfg_path = _dj_config_path(project_dir)
-    if not os.path.exists(dj_cfg_path):
+    if not Path(dj_cfg_path).exists():
         raise FileNotFoundError(
             f"Missing DJ config at {dj_cfg_path}. Create it with dj_generate_config()."
         )
@@ -135,7 +134,7 @@ def update_kpms_dj_config(project_dir: str, **kwargs) -> Dict[str, Any]:
     keypoint_moseq.io.update_config), then rewrite the file and return the dict.
     """
     dj_cfg_path = _dj_config_path(project_dir)
-    if not os.path.exists(dj_cfg_path):
+    if not Path(dj_cfg_path).exists():
         raise FileNotFoundError(
             f"Missing DJ config at {dj_cfg_path}. Create it with dj_generate_config()."
         )

element_moseq/version.py

@@ -2,4 +2,4 @@
 Package metadata
 """
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "element-moseq"
-version = "1.0.0"
+version = "1.0.1"
 description = "Keypoint-MoSeq DataJoint Element"
 readme = "README.md"
 license = {text = "MIT"}
@@ -25,6 +25,7 @@ dependencies = [
     "ipywidgets",
     "opencv-python",
     "keypoint-moseq @ git+https://github.com/dattalab/keypoint-moseq/",
+    "pdf2image",
 ]
 
 [project.optional-dependencies]