config.py

"""
Central configuration for the de-identification pipeline.
Edit only this file to change behaviour — nothing else needs to be touched.
"""

from pathlib import Path

ROOT = Path(__file__).parent

# ── Folder layout ─────────────────────────────────────────────────────────────
DATASET_DIR       = ROOT / "dataset"                  # raw images + labels
DATASET_YOLO      = ROOT / "dataset_yolo"             # YOLO dataset (baseline)
DATASET_CLAHE     = ROOT / "dataset_yolo_clahe"       # YOLO dataset (CLAHE+Canny)
MODELS_DIR        = ROOT / "models"
OUTPUTS_DIR       = ROOT / "outputs"

YOLO_BASELINE_W   = MODELS_DIR / "yolo"        / "baseline"    / "weights" / "best.pt"
YOLO_CLAHE_W      = MODELS_DIR / "yolo"        / "clahe_canny" / "weights" / "best.pt"
INPAINT_SSL_W     = MODELS_DIR / "inpainting"  / "ssl"         / "best.pt"

# ── Pipeline flags ─────────────────────────────────────────────────────────────
# Which detector to use:  "baseline"  |  "clahe_canny"
DETECTOR = "clahe_canny"

# False → load existing weights (default).  True → retrain from scratch.
RETRAIN_YOLO = False

# Anonymization method:  "blur"  |  "lama"  |  "ssl"
ANONYMIZER = "lama"

# Only relevant when ANONYMIZER == "ssl"
RETRAIN_INPAINT = True

# OCR engine:  "paddle" (default)  |  "tesseract"  |  "easyocr"
OCR_ENGINE = "paddle"

# ── YOLO hyperparameters ──────────────────────────────────────────────────────
YOLO_EPOCHS = 50
YOLO_BATCH  = 8
YOLO_IMGSZ  = 640

# ── Inpainting hyperparameters ────────────────────────────────────────────────
INPAINT_PATCH        = 128
INPAINT_EPOCHS       = 50
INPAINT_BATCH        = 32
INPAINT_LR           = 3e-4
INPAINT_MASK_H_SCALE = 3.0   # inflate mask height during training to avoid gray collapse
INPAINT_N_PATCH      = 30    # patches sampled per training image

# ── Inference ─────────────────────────────────────────────────────────────────
CONF        = 0.25
BLUR_KERNEL = 51
BLUR_PASSES = 3
MASK_OP     = "dilate"   # "dilate" → expand box  |  "erode" → shrink box
DILATE_PX   = 1         # pixels added around each detected box before inpainting
ERODE_PX    = 3          # pixels removed from each detected box before inpainting

CLASS_NAMES  = ["name", "id", "age", "date", "time"]
CLASS_COLORS = {       # BGR
    0: (80,  80,  255),   # name  → red
    1: (255, 200,  80),   # id    → cyan
    2: (80,  255,  80),   # age   → green
    3: (80,  200, 255),   # date  → orange
    4: (255,  80, 200),   # time  → purple
}