config/config.py

"""
Configuration file for the OCR Number Plate Pipeline
"""
import os
from pathlib import Path

# ============== PATHS ==============
BASE_DIR = Path(__file__).parent.parent
PROJECT_ROOT = BASE_DIR.parent

# Model paths
YOLO_MODEL_PATH = Path("final_models") / "best12345.pt"
OCR_MODEL_PATH = Path("final_models") / "ocr_model.pth"
LABEL_MAP_PATH = Path("label_map.json")

# Output paths
OUTPUT_DIR = PROJECT_ROOT / "output"
CONTOURS_DIR = OUTPUT_DIR / "contours"
CONTOURS_BW_DIR = OUTPUT_DIR / "contours_bw"
RESULTS_DIR = OUTPUT_DIR / "results"

# ============== OCR MODEL CONFIG ==============
OCR_CONFIG = {
    "input_size": (128, 128),
    "num_classes": 71,  # 0-9, A-Z, Nepali chars, Nepali Flag
    "backbone": "resnet18",
    "pretrained": False,
}

# ============== YOLO CONFIG ==============
YOLO_CONFIG = {
    "confidence_threshold": 0.5,
    "iou_threshold": 0.45,
    "img_size": 640,
    "device": "auto",  # "auto", "cuda", "cpu"
}

# ============== PREPROCESSING CONFIG ==============
PREPROCESS_CONFIG = {
    "clahe_clip_limit": 2.0,
    "clahe_grid_size": (8, 8),
    "gaussian_blur_kernel": (3, 3),
    "binary_threshold": 180,
    "otsu_threshold": True,
}

# ============== CONTOUR DETECTION CONFIG ==============
CONTOUR_CONFIG = {
    "min_area": 60,
    "min_width": 3,
    "min_height": 2,
    "min_aspect_ratio": 0.08,
    "max_aspect_ratio": 4.0,
    "max_width_ratio": 0.55,
    "max_height_ratio": 0.95,
    "min_height_rel_median": 0.45,
    "max_height_rel_median": 2.2,
    "min_width_rel_median": 0.35,
    "max_width_rel_median": 2.4,
    "min_area_rel_median": 0.20,
    "max_area_rel_median": 3.8,
    "padding": 1,
    "center_threshold": 15,  # For removing overlapping contours
    "skip_borders": 0,  # Legacy key (no blind skipping)
    "remove_edge_artifacts": True,
    "edge_margin": 2,
}

# ============== LINE GROUPING CONFIG ==============
LINE_CONFIG = {
    "y_threshold_ratio": 0.5,  # Ratio of average height for line grouping
    "default_y_threshold": 20,
}

# ============== OCR INFERENCE CONFIG ==============
INFERENCE_CONFIG = {
    "min_confidence": 0.10,  # 10% minimum confidence
    "batch_size": 16,
}

# ============== VISUALIZATION CONFIG ==============
VIZ_CONFIG = {
    "show_plots": True,
    "save_results": True,
    "figure_size": (12, 6),
    "font_size": 9,
    "max_cols": 10,
}

# ============== DEVICE CONFIGURATION ==============
def get_device():
    """Auto-detect the best available device."""
    import torch
    if YOLO_CONFIG["device"] == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(YOLO_CONFIG["device"])

# ============== CREATE DIRECTORIES ==============
def setup_directories():
    """Create necessary output directories."""
    for dir_path in [OUTPUT_DIR, CONTOURS_DIR, CONTOURS_BW_DIR, RESULTS_DIR]:
        dir_path.mkdir(parents=True, exist_ok=True)