#!/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())