orientation/orient.py

#!/usr/bin/env python3
"""Apply the fixed PP-LCNet document-orientation policy on CPU."""

from __future__ import annotations

import argparse
import json
import os
import shutil
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import cv2
import numpy as np
import paddle.inference as paddle_infer
from PIL import Image, ImageOps


CLASSIFIER = "PP-LCNet_x1_0_doc_ori"
THRESHOLD = 0.85
LABELS = (0, 90, 180, 270)
EXTENSIONS = {".jpg", ".jpeg", ".png"}
MEAN = np.asarray((0.485, 0.456, 0.406), dtype=np.float32)
STD = np.asarray((0.229, 0.224, 0.225), dtype=np.float32)
TRANSPOSE = {
    90: Image.Transpose.ROTATE_90,
    180: Image.Transpose.ROTATE_180,
    270: Image.Transpose.ROTATE_270,
}
MODEL_FILES = ("inference.json", "inference.pdiparams")
MATERIALIZE_WORKERS = 8


def verify_model_files(model_dir: Path) -> None:
    missing = [name for name in MODEL_FILES if not (model_dir / name).is_file()]
    if missing:
        raise RuntimeError(f"orientation model files are missing: {missing}")


def preprocess(path: Path) -> np.ndarray:
    with Image.open(path) as opened:
        image = np.asarray(ImageOps.exif_transpose(opened).convert("RGB"))
    height, width = image.shape[:2]
    scale = 256.0 / min(height, width)
    resized = cv2.resize(
        image,
        (round(width * scale), round(height * scale)),
        interpolation=cv2.INTER_LINEAR,
    )
    height, width = resized.shape[:2]
    top = (height - 224) // 2
    left = (width - 224) // 2
    crop = resized[top : top + 224, left : left + 224].astype(np.float32)
    crop = (crop / 255.0 - MEAN) / STD
    return np.transpose(crop, (2, 0, 1))


class Predictor:
    def __init__(self, model_dir: Path):
        config = paddle_infer.Config(
            str(model_dir / "inference.json"),
            str(model_dir / "inference.pdiparams"),
        )
        config.disable_gpu()
        config.set_cpu_math_library_num_threads(4)
        config.disable_glog_info()
        self.predictor = paddle_infer.create_predictor(config)
        self.input_name = self.predictor.get_input_names()[0]
        self.output_name = self.predictor.get_output_names()[0]

    def predict(self, paths: list[Path]) -> list[tuple[Path, int, float]]:
        rows: list[tuple[Path, int, float]] = []
        with ThreadPoolExecutor(max_workers=4) as pool:
            for start in range(0, len(paths), 8):
                batch_paths = paths[start : start + 8]
                batch = np.stack(list(pool.map(preprocess, batch_paths)))
                input_handle = self.predictor.get_input_handle(self.input_name)
                input_handle.reshape(batch.shape)
                input_handle.copy_from_cpu(batch)
                self.predictor.run()
                values = self.predictor.get_output_handle(self.output_name).copy_to_cpu()
                if values.ndim != 2 or values.shape[1] != len(LABELS):
                    raise RuntimeError(
                        f"unexpected classifier output shape: {values.shape}"
                    )
                if np.any(values < 0) or not np.allclose(
                    values.sum(axis=1), 1.0, atol=1e-3
                ):
                    values = np.exp(values - values.max(axis=1, keepdims=True))
                    values /= values.sum(axis=1, keepdims=True)
                for path, probabilities in zip(batch_paths, values, strict=True):
                    index = int(np.argmax(probabilities))
                    rows.append((path, LABELS[index], float(probabilities[index])))
        return rows


def materialize(
    source: Path, angle: int, confidence: float, output_dir: Path
) -> bool:
    selected = angle in TRANSPOSE and confidence >= THRESHOLD
    if selected:
        target = output_dir / f"{source.stem}.png"
        with Image.open(source) as opened:
            image = ImageOps.exif_transpose(opened).convert("RGB")
            image.transpose(TRANSPOSE[angle]).save(
                target, format="PNG", optimize=False
            )
    else:
        target = output_dir / source.name
        try:
            os.link(source, target)
        except OSError:
            shutil.copy2(source, target)
    return selected


def main() -> int:
    started = time.perf_counter()
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input-dir", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--model-dir", required=True, type=Path)
    parser.add_argument(
        "--record-path",
        type=Path,
        default=None,
        help="Optional JSONL file recording label/confidence/rotated per image.",
    )
    args = parser.parse_args()

    if not args.input_dir.is_dir():
        parser.error(f"input directory not found: {args.input_dir}")
    if not args.model_dir.is_dir():
        parser.error(f"orientation model directory not found: {args.model_dir}")
    if args.input_dir.resolve() == args.output_dir.resolve():
        parser.error("input and output directories must differ")
    if args.output_dir.is_dir() and any(args.output_dir.iterdir()):
        parser.error(f"output directory must be empty: {args.output_dir}")

    paths = sorted(
        path
        for path in args.input_dir.iterdir()
        if path.is_file() and path.suffix.lower() in EXTENSIONS
    )
    if not paths:
        parser.error(f"no supported images found in: {args.input_dir}")
    duplicate_stems = [
        stem for stem, count in Counter(path.stem for path in paths).items() if count > 1
    ]
    if duplicate_stems:
        raise RuntimeError(f"duplicate input stems: {sorted(duplicate_stems)}")

    verify_model_files(args.model_dir)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    predictions = Predictor(args.model_dir).predict(paths)
    classified = time.perf_counter()
    with ThreadPoolExecutor(max_workers=MATERIALIZE_WORKERS) as pool:
        futures = [
            pool.submit(materialize, source, angle, confidence, args.output_dir)
            for source, angle, confidence in predictions
        ]
        results = [future.result() for future in futures]
        selected = sum(results)
    if args.record_path is not None:
        args.record_path.parent.mkdir(parents=True, exist_ok=True)
        with args.record_path.open("w", encoding="utf-8") as handle:
            for (source, angle, confidence), rotated in zip(predictions, results, strict=True):
                handle.write(
                    json.dumps(
                        {
                            "image": source.name,
                            "label": angle,
                            "confidence": round(float(confidence), 6),
                            "rotated": bool(rotated),
                        }
                    )
                    + "\n"
                )
    finished = time.perf_counter()
    print(
        f"[orientation] classifier={CLASSIFIER} images={len(paths)} "
        f"corrected={selected} workers={MATERIALIZE_WORKERS} "
        f"classify_s={classified - started:.1f} "
        f"materialize_s={finished - classified:.1f} "
        f"total_s={finished - started:.1f} "
        f"images_per_s={len(paths) / (finished - started):.2f}"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())