landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/audit.py

@@ -0,0 +1,338 @@
+"""Clinical audit report generator for regulatory compliance.
+
+Generates structured HTML reports summarizing safety validation results,
+model performance, and Fitzpatrick equity analysis for clinical review.
+
+Reports include:
+- Safety validation pass/fail summary per patient
+- Aggregate statistics by procedure and Fitzpatrick type
+- Flagged cases for manual review
+- Model version and configuration provenance
+
+Usage:
+    from landmarkdiff.audit import AuditReporter, AuditCase
+
+    reporter = AuditReporter(model_version="0.3.2")
+    reporter.add_case(AuditCase(
+        case_id="P001",
+        procedure="rhinoplasty",
+        safety_passed=True,
+        identity_sim=0.87,
+        fitzpatrick_type="III",
+    ))
+    reporter.generate_report("audit_report.html")
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class AuditCase:
+    """A single patient case for audit reporting."""
+
+    case_id: str
+    procedure: str
+    safety_passed: bool
+    identity_sim: float = 0.0
+    intensity: float = 65.0
+    fitzpatrick_type: str = ""
+    warnings: list[str] = field(default_factory=list)
+    failures: list[str] = field(default_factory=list)
+    metrics: dict[str, float] = field(default_factory=dict)
+    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+
+@dataclass
+class AuditSummary:
+    """Aggregate statistics for an audit report."""
+
+    total_cases: int = 0
+    passed_cases: int = 0
+    failed_cases: int = 0
+    flagged_cases: int = 0
+    pass_rate: float = 0.0
+    mean_identity_sim: float = 0.0
+    by_procedure: dict[str, dict[str, Any]] = field(default_factory=dict)
+    by_fitzpatrick: dict[str, dict[str, Any]] = field(default_factory=dict)
+
+
+class AuditReporter:
+    """Generate clinical audit reports from safety validation results.
+
+    Args:
+        model_version: Model version string for provenance.
+        report_title: Title for generated reports.
+    """
+
+    def __init__(
+        self,
+        model_version: str = "0.3.2",
+        report_title: str = "LandmarkDiff Clinical Audit Report",
+    ) -> None:
+        self.model_version = model_version
+        self.report_title = report_title
+        self.cases: list[AuditCase] = []
+
+    def add_case(self, case: AuditCase) -> None:
+        """Add a case to the audit report."""
+        self.cases.append(case)
+
+    def add_cases(self, cases: list[AuditCase]) -> None:
+        """Add multiple cases."""
+        self.cases.extend(cases)
+
+    def clear(self) -> None:
+        """Clear all cases."""
+        self.cases.clear()
+
+    def compute_summary(self) -> AuditSummary:
+        """Compute aggregate statistics from all cases."""
+        if not self.cases:
+            return AuditSummary()
+
+        total = len(self.cases)
+        passed = sum(1 for c in self.cases if c.safety_passed)
+        failed = total - passed
+        flagged = sum(1 for c in self.cases if not c.safety_passed or c.warnings)
+
+        id_sims = [c.identity_sim for c in self.cases if c.identity_sim > 0]
+        mean_id = sum(id_sims) / len(id_sims) if id_sims else 0.0
+
+        # By procedure
+        by_proc: dict[str, dict[str, Any]] = {}
+        for case in self.cases:
+            proc = case.procedure
+            if proc not in by_proc:
+                by_proc[proc] = {"total": 0, "passed": 0, "id_sims": []}
+            by_proc[proc]["total"] += 1
+            if case.safety_passed:
+                by_proc[proc]["passed"] += 1
+            if case.identity_sim > 0:
+                by_proc[proc]["id_sims"].append(case.identity_sim)
+
+        for _proc, stats in by_proc.items():
+            stats["pass_rate"] = stats["passed"] / max(stats["total"], 1)
+            stats["mean_identity_sim"] = (
+                sum(stats["id_sims"]) / len(stats["id_sims"]) if stats["id_sims"] else 0.0
+            )
+            del stats["id_sims"]
+
+        # By Fitzpatrick type
+        by_fitz: dict[str, dict[str, Any]] = {}
+        for case in self.cases:
+            ft = case.fitzpatrick_type or "Unknown"
+            if ft not in by_fitz:
+                by_fitz[ft] = {"total": 0, "passed": 0, "id_sims": []}
+            by_fitz[ft]["total"] += 1
+            if case.safety_passed:
+                by_fitz[ft]["passed"] += 1
+            if case.identity_sim > 0:
+                by_fitz[ft]["id_sims"].append(case.identity_sim)
+
+        for _ft, stats in by_fitz.items():
+            stats["pass_rate"] = stats["passed"] / max(stats["total"], 1)
+            stats["mean_identity_sim"] = (
+                sum(stats["id_sims"]) / len(stats["id_sims"]) if stats["id_sims"] else 0.0
+            )
+            del stats["id_sims"]
+
+        return AuditSummary(
+            total_cases=total,
+            passed_cases=passed,
+            failed_cases=failed,
+            flagged_cases=flagged,
+            pass_rate=passed / total,
+            mean_identity_sim=mean_id,
+            by_procedure=by_proc,
+            by_fitzpatrick=by_fitz,
+        )
+
+    def flagged_cases(self) -> list[AuditCase]:
+        """Return cases that need manual review (failed or have warnings)."""
+        return [c for c in self.cases if not c.safety_passed or c.warnings]
+
+    def to_json(self) -> str:
+        """Export audit data as JSON."""
+        summary = self.compute_summary()
+        data = {
+            "report_title": self.report_title,
+            "model_version": self.model_version,
+            "generated_at": datetime.now(timezone.utc).isoformat(),
+            "summary": {
+                "total_cases": summary.total_cases,
+                "passed_cases": summary.passed_cases,
+                "failed_cases": summary.failed_cases,
+                "flagged_cases": summary.flagged_cases,
+                "pass_rate": round(summary.pass_rate, 4),
+                "mean_identity_sim": round(summary.mean_identity_sim, 4),
+            },
+            "by_procedure": {
+                k: {kk: round(vv, 4) if isinstance(vv, float) else vv for kk, vv in v.items()}
+                for k, v in summary.by_procedure.items()
+            },
+            "by_fitzpatrick": {
+                k: {kk: round(vv, 4) if isinstance(vv, float) else vv for kk, vv in v.items()}
+                for k, v in summary.by_fitzpatrick.items()
+            },
+            "cases": [
+                {
+                    "case_id": c.case_id,
+                    "procedure": c.procedure,
+                    "safety_passed": c.safety_passed,
+                    "identity_sim": round(c.identity_sim, 4),
+                    "intensity": c.intensity,
+                    "fitzpatrick_type": c.fitzpatrick_type,
+                    "warnings": c.warnings,
+                    "failures": c.failures,
+                    "metrics": {k: round(v, 4) for k, v in c.metrics.items()},
+                    "timestamp": c.timestamp,
+                }
+                for c in self.cases
+            ],
+        }
+        return json.dumps(data, indent=2)
+
+    def generate_report(self, output_path: str | Path) -> Path:
+        """Generate an HTML audit report.
+
+        Args:
+            output_path: Path to save the HTML report.
+
+        Returns:
+            Path to the generated report.
+        """
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        summary = self.compute_summary()
+        html = self._render_html(summary)
+
+        output_path.write_text(html)
+        return output_path
+
+    def _render_html(self, summary: AuditSummary) -> str:
+        """Render the audit report as HTML."""
+        now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
+        status = "PASS" if summary.failed_cases == 0 else "REQUIRES REVIEW"
+        status_color = "#28a745" if summary.failed_cases == 0 else "#dc3545"
+
+        # Build procedure rows
+        proc_rows = ""
+        for proc, stats in sorted(summary.by_procedure.items()):
+            rate = stats["pass_rate"]
+            rate_color = "#28a745" if rate >= 0.95 else "#ffc107" if rate >= 0.8 else "#dc3545"
+            proc_rows += (
+                f"<tr>"
+                f"<td>{proc.title()}</td>"
+                f"<td>{stats['total']}</td>"
+                f"<td>{stats['passed']}</td>"
+                f'<td style="color:{rate_color};font-weight:bold">{rate:.1%}</td>'
+                f"<td>{stats['mean_identity_sim']:.4f}</td>"
+                f"</tr>\n"
+            )
+
+        # Build Fitzpatrick rows
+        fitz_rows = ""
+        for ft, stats in sorted(summary.by_fitzpatrick.items()):
+            rate = stats["pass_rate"]
+            rate_color = "#28a745" if rate >= 0.95 else "#ffc107" if rate >= 0.8 else "#dc3545"
+            fitz_rows += (
+                f"<tr>"
+                f"<td>{ft}</td>"
+                f"<td>{stats['total']}</td>"
+                f"<td>{stats['passed']}</td>"
+                f'<td style="color:{rate_color};font-weight:bold">{rate:.1%}</td>'
+                f"<td>{stats['mean_identity_sim']:.4f}</td>"
+                f"</tr>\n"
+            )
+
+        # Build flagged cases
+        flagged = self.flagged_cases()
+        flagged_rows = ""
+        for c in flagged:
+            issues = "; ".join(c.failures + [f"WARN: {w}" for w in c.warnings])
+            bg = "#fff3cd" if c.safety_passed else "#f8d7da"
+            flagged_rows += (
+                f'<tr style="background:{bg}">'
+                f"<td>{c.case_id}</td>"
+                f"<td>{c.procedure.title()}</td>"
+                f"<td>{c.fitzpatrick_type}</td>"
+                f"<td>{c.identity_sim:.4f}</td>"
+                f"<td>{'WARN' if c.safety_passed else 'FAIL'}</td>"
+                f"<td>{issues}</td>"
+                f"</tr>\n"
+            )
+
+        return f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<title>{self.report_title}</title>
+<style>
+body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
+       max-width: 1100px; margin: 0 auto; padding: 20px; color: #333; }}
+h1 {{ border-bottom: 3px solid #333; padding-bottom: 10px; }}
+h2 {{ color: #555; margin-top: 30px; border-bottom: 1px solid #ddd; padding-bottom: 5px; }}
+table {{ border-collapse: collapse; width: 100%; margin: 15px 0; }}
+th, td {{ border: 1px solid #ddd; padding: 8px 12px; text-align: left; }}
+th {{ background: #f8f9fa; font-weight: 600; }}
+tr:hover {{ background: #f5f5f5; }}
+.status {{ display: inline-block; padding: 4px 12px; border-radius: 4px;
+           color: white; font-weight: bold; font-size: 18px; }}
+.summary-grid {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(180px, 1fr));
+                 gap: 15px; margin: 20px 0; }}
+.summary-card {{ background: #f8f9fa; border-radius: 8px; padding: 15px; text-align: center; }}
+.summary-card .value {{ font-size: 28px; font-weight: bold; color: #333; }}
+.summary-card .label {{ font-size: 12px; color: #888; text-transform: uppercase; }}
+.disclaimer {{ background: #fff3cd; border: 1px solid #ffc107; border-radius: 4px;
+               padding: 12px; margin: 20px 0; font-size: 13px; }}
+footer {{ margin-top: 40px; padding-top: 15px; border-top: 1px solid #ddd;
+          font-size: 12px; color: #999; }}
+</style>
+</head>
+<body>
+<h1>{self.report_title}</h1>
+<p>Generated: {now} &nbsp;|&nbsp; Model version: <code>{self.model_version}</code>
+&nbsp;|&nbsp; Overall status: <span class="status" style="background:{status_color}">{status}</span></p>
+
+<div class="disclaimer">
+<strong>Disclaimer:</strong> This report is for research and development purposes only.
+LandmarkDiff predictions are AI-generated visualizations and do not constitute medical advice
+or guarantee surgical outcomes. All predictions should be reviewed by qualified clinical professionals.
+</div>
+
+<h2>Summary</h2>
+<div class="summary-grid">
+<div class="summary-card"><div class="value">{summary.total_cases}</div><div class="label">Total Cases</div></div>
+<div class="summary-card"><div class="value" style="color:#28a745">{summary.passed_cases}</div><div class="label">Passed</div></div>
+<div class="summary-card"><div class="value" style="color:#dc3545">{summary.failed_cases}</div><div class="label">Failed</div></div>
+<div class="summary-card"><div class="value" style="color:#ffc107">{summary.flagged_cases}</div><div class="label">Flagged</div></div>
+<div class="summary-card"><div class="value">{summary.pass_rate:.1%}</div><div class="label">Pass Rate</div></div>
+<div class="summary-card"><div class="value">{summary.mean_identity_sim:.4f}</div><div class="label">Mean ID Sim</div></div>
+</div>
+
+<h2>Performance by Procedure</h2>
+<table>
+<tr><th>Procedure</th><th>Total</th><th>Passed</th><th>Pass Rate</th><th>Mean ID Sim</th></tr>
+{proc_rows}</table>
+
+<h2>Equity Analysis by Fitzpatrick Type</h2>
+<table>
+<tr><th>Fitzpatrick Type</th><th>Total</th><th>Passed</th><th>Pass Rate</th><th>Mean ID Sim</th></tr>
+{fitz_rows}</table>
+
+{"<h2>Flagged Cases (Require Review)</h2>" if flagged_rows else ""}
+{"<table><tr><th>Case ID</th><th>Procedure</th><th>Fitzpatrick</th><th>ID Sim</th><th>Status</th><th>Issues</th></tr>" + flagged_rows + "</table>" if flagged_rows else "<p>No flagged cases.</p>"}
+
+<footer>
+LandmarkDiff v{self.model_version} &mdash; Clinical Audit Report &mdash;
+For research use only. Not FDA approved.
+</footer>
+</body>
+</html>"""

landmarkdiff/augmentation.py

@@ -0,0 +1,293 @@
+"""Training data augmentation pipeline for LandmarkDiff.
+
+Provides domain-specific augmentations that maintain landmark consistency:
+- Geometric: flip, rotation, affine (landmarks co-transformed)
+- Photometric: color jitter, brightness, contrast (applied to images only)
+- Skin-tone augmentation: ITA-space perturbation for Fitzpatrick balance
+- Conditioning augmentation: noise injection, dropout for robustness
+
+All augmentations preserve the correspondence between:
+  input_image ↔ conditioning_image ↔ target_image ↔ mask
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+
+
+@dataclass
+class AugmentationConfig:
+    """Augmentation parameters."""
+
+    # Geometric
+    random_flip: bool = True
+    random_rotation_deg: float = 5.0
+    random_scale: tuple[float, float] = (0.95, 1.05)
+    random_translate: float = 0.02  # fraction of image size
+
+    # Photometric (images only, not conditioning)
+    brightness_range: tuple[float, float] = (0.9, 1.1)
+    contrast_range: tuple[float, float] = (0.9, 1.1)
+    saturation_range: tuple[float, float] = (0.9, 1.1)
+    hue_shift_range: float = 5.0  # degrees
+
+    # Conditioning augmentation
+    conditioning_dropout_prob: float = 0.1
+    conditioning_noise_std: float = 0.02
+
+    # Skin-tone augmentation
+    ita_perturbation_std: float = 3.0  # ITA angle noise
+
+    seed: int | None = None
+
+
+def augment_training_sample(
+    input_image: np.ndarray,
+    target_image: np.ndarray,
+    conditioning: np.ndarray,
+    mask: np.ndarray,
+    landmarks_src: np.ndarray | None = None,
+    landmarks_dst: np.ndarray | None = None,
+    config: AugmentationConfig | None = None,
+    rng: np.random.Generator | None = None,
+) -> dict[str, np.ndarray]:
+    """Apply consistent augmentations to a training sample.
+
+    All spatial transforms are applied to images AND landmarks together
+    so correspondence is preserved.
+
+    Args:
+        input_image: (H, W, 3) original face image (uint8 BGR).
+        target_image: (H, W, 3) target face image (uint8 BGR).
+        conditioning: (H, W, 3) conditioning image (uint8).
+        mask: (H, W) or (H, W, 1) float32 mask.
+        landmarks_src: (N, 2) normalized [0,1] source landmark coords.
+        landmarks_dst: (N, 2) normalized [0,1] target landmark coords.
+        config: Augmentation parameters.
+        rng: Random generator for reproducibility.
+
+    Returns:
+        Dict with augmented versions of all inputs.
+    """
+    if config is None:
+        config = AugmentationConfig()
+    if rng is None:
+        rng = np.random.default_rng(config.seed)
+
+    h, w = input_image.shape[:2]
+    out_input = input_image.copy()
+    out_target = target_image.copy()
+    out_cond = conditioning.copy()
+    out_mask = mask.copy()
+    out_lm_src = landmarks_src.copy() if landmarks_src is not None else None
+    out_lm_dst = landmarks_dst.copy() if landmarks_dst is not None else None
+
+    # --- Geometric augmentations (applied to all) ---
+
+    # Random horizontal flip
+    if config.random_flip and rng.random() < 0.5:
+        out_input = np.ascontiguousarray(out_input[:, ::-1])
+        out_target = np.ascontiguousarray(out_target[:, ::-1])
+        out_cond = np.ascontiguousarray(out_cond[:, ::-1])
+        out_mask = np.ascontiguousarray(
+            out_mask[:, ::-1] if out_mask.ndim == 2 else out_mask[:, ::-1, :]
+        )
+        if out_lm_src is not None:
+            out_lm_src[:, 0] = 1.0 - out_lm_src[:, 0]
+        if out_lm_dst is not None:
+            out_lm_dst[:, 0] = 1.0 - out_lm_dst[:, 0]
+
+    # Random rotation + scale + translate
+    if config.random_rotation_deg > 0 or config.random_scale != (1.0, 1.0):
+        angle = rng.uniform(-config.random_rotation_deg, config.random_rotation_deg)
+        scale = rng.uniform(config.random_scale[0], config.random_scale[1])
+        tx = rng.uniform(-config.random_translate, config.random_translate) * w
+        ty = rng.uniform(-config.random_translate, config.random_translate) * h
+
+        center = (w / 2, h / 2)
+        M = cv2.getRotationMatrix2D(center, angle, scale)
+        M[0, 2] += tx
+        M[1, 2] += ty
+
+        out_input = cv2.warpAffine(out_input, M, (w, h), borderMode=cv2.BORDER_REFLECT_101)
+        out_target = cv2.warpAffine(out_target, M, (w, h), borderMode=cv2.BORDER_REFLECT_101)
+        out_cond = cv2.warpAffine(
+            out_cond, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=0
+        )
+        mask_2d = out_mask if out_mask.ndim == 2 else out_mask[:, :, 0]
+        mask_2d = cv2.warpAffine(mask_2d, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=0)
+        out_mask = mask_2d if out_mask.ndim == 2 else mask_2d[:, :, np.newaxis]
+
+        # Transform landmarks
+        if out_lm_src is not None:
+            out_lm_src = _transform_landmarks(out_lm_src, M, w, h)
+        if out_lm_dst is not None:
+            out_lm_dst = _transform_landmarks(out_lm_dst, M, w, h)
+
+    # --- Photometric augmentations (images only, not conditioning/mask) ---
+
+    # Brightness
+    b_factor = rng.uniform(config.brightness_range[0], config.brightness_range[1])
+    out_input = np.clip(out_input.astype(np.float32) * b_factor, 0, 255).astype(np.uint8)
+    out_target = np.clip(out_target.astype(np.float32) * b_factor, 0, 255).astype(np.uint8)
+
+    # Contrast
+    c_factor = rng.uniform(config.contrast_range[0], config.contrast_range[1])
+    mean_in = out_input.mean()
+    mean_tgt = out_target.mean()
+    out_input = np.clip(
+        (out_input.astype(np.float32) - mean_in) * c_factor + mean_in, 0, 255
+    ).astype(np.uint8)
+    out_target = np.clip(
+        (out_target.astype(np.float32) - mean_tgt) * c_factor + mean_tgt, 0, 255
+    ).astype(np.uint8)
+
+    # Saturation (in HSV space)
+    s_factor = rng.uniform(config.saturation_range[0], config.saturation_range[1])
+    if abs(s_factor - 1.0) > 1e-4:
+        out_input = _adjust_saturation(out_input, s_factor)
+        out_target = _adjust_saturation(out_target, s_factor)
+
+    # Hue shift
+    if config.hue_shift_range > 0:
+        hue_delta = rng.uniform(-config.hue_shift_range, config.hue_shift_range)
+        if abs(hue_delta) > 0.1:
+            out_input = _shift_hue(out_input, hue_delta)
+            out_target = _shift_hue(out_target, hue_delta)
+
+    # --- Conditioning augmentation ---
+
+    # Conditioning dropout (replace with zeros to learn unconditional)
+    if config.conditioning_dropout_prob > 0 and rng.random() < config.conditioning_dropout_prob:
+        out_cond = np.zeros_like(out_cond)
+
+    # Conditioning noise
+    if config.conditioning_noise_std > 0:
+        noise = rng.normal(0, config.conditioning_noise_std * 255, out_cond.shape)
+        out_cond = np.clip(out_cond.astype(np.float32) + noise, 0, 255).astype(np.uint8)
+
+    result = {
+        "input_image": out_input,
+        "target_image": out_target,
+        "conditioning": out_cond,
+        "mask": out_mask,
+    }
+    if out_lm_src is not None:
+        result["landmarks_src"] = out_lm_src
+    if out_lm_dst is not None:
+        result["landmarks_dst"] = out_lm_dst
+
+    return result
+
+
+def _transform_landmarks(landmarks: np.ndarray, M: np.ndarray, w: int, h: int) -> np.ndarray:
+    """Transform normalized landmarks with an affine matrix."""
+    # Convert to pixel coords
+    px = landmarks.copy()
+    px[:, 0] *= w
+    px[:, 1] *= h
+
+    # Apply affine transform
+    ones = np.ones((px.shape[0], 1))
+    px_h = np.hstack([px, ones])  # (N, 3)
+    transformed = (M @ px_h.T).T  # (N, 2)
+
+    # Back to normalized
+    transformed[:, 0] /= w
+    transformed[:, 1] /= h
+    return np.clip(transformed, 0.0, 1.0)
+
+
+def _adjust_saturation(img: np.ndarray, factor: float) -> np.ndarray:
+    """Adjust saturation of a BGR image."""
+    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float32)
+    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * factor, 0, 255)
+    return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
+
+
+def _shift_hue(img: np.ndarray, delta_deg: float) -> np.ndarray:
+    """Shift hue of a BGR image by delta degrees."""
+    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float32)
+    # OpenCV hue range is [0, 180]
+    hsv[:, :, 0] = (hsv[:, :, 0] + delta_deg / 2) % 180
+    return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
+
+
+def augment_skin_tone(
+    image: np.ndarray,
+    ita_delta: float = 0.0,
+) -> np.ndarray:
+    """Augment skin tone by shifting in L*a*b* space.
+
+    This helps balance Fitzpatrick representation in training by
+    simulating different skin tones from existing samples.
+
+    Args:
+        image: (H, W, 3) BGR uint8 image.
+        ita_delta: ITA angle shift (positive = lighter, negative = darker).
+
+    Returns:
+        Augmented image with shifted skin tone.
+    """
+    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB).astype(np.float32)
+
+    # Shift L channel (lightness) based on ITA delta
+    # ITA = arctan((L-50)/b), so shifting ITA shifts L
+    l_shift = ita_delta * 0.5  # approximate mapping
+    lab[:, :, 0] = np.clip(lab[:, :, 0] + l_shift, 0, 255)
+
+    # Slightly shift b channel too for more natural tone changes
+    b_shift = -ita_delta * 0.15
+    lab[:, :, 2] = np.clip(lab[:, :, 2] + b_shift, 0, 255)
+
+    return cv2.cvtColor(lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
+
+
+class FitzpatrickBalancer:
+    """Oversample underrepresented Fitzpatrick types during training.
+
+    Maintains per-type counts and generates sampling weights to ensure
+    equitable training across all skin types.
+    """
+
+    def __init__(self, target_distribution: dict[str, float] | None = None):
+        """Initialize balancer.
+
+        Args:
+            target_distribution: Target fraction per type. Defaults to uniform.
+        """
+        self.target = target_distribution or {
+            "I": 1 / 6,
+            "II": 1 / 6,
+            "III": 1 / 6,
+            "IV": 1 / 6,
+            "V": 1 / 6,
+            "VI": 1 / 6,
+        }
+        self._counts: dict[str, int] = {}
+
+    def register_sample(self, fitz_type: str) -> None:
+        """Register a sample's Fitzpatrick type."""
+        self._counts[fitz_type] = self._counts.get(fitz_type, 0) + 1
+
+    def get_sampling_weights(self, fitz_types: list[str]) -> np.ndarray:
+        """Compute sampling weights for a list of samples.
+
+        Returns weights inversely proportional to type frequency,
+        so underrepresented types get upsampled.
+        """
+        total = sum(self._counts.values()) or 1
+        weights = []
+        for ft in fitz_types:
+            count = self._counts.get(ft, 1)
+            freq = count / total
+            target_freq = self.target.get(ft, 1 / 6)
+            # Weight = target / actual (capped for stability)
+            w = min(target_freq / max(freq, 1e-6), 5.0)
+            weights.append(w)
+
+        w = np.array(weights, dtype=np.float64)
+        return w / w.sum()  # normalize to probability distribution