File size: 1,814 Bytes
b8d39bc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """
ROI detection: full passport image -> bounding box of the dotted series/number strip.
Uses a YOLO model fine-tuned on Ukrainian passport images (class: "Dotted").
"""
from pathlib import Path
import cv2
import numpy as np
def load_detector(model_path: str | Path):
from ultralytics import YOLO
return YOLO(str(model_path))
def detect_roi(image: np.ndarray, model, conf: float = 0.30) -> tuple[int, int, int, int] | None:
"""
Detect the dotted passport series/number strip in a full passport image.
Args:
image: BGR image (numpy array) or grayscale — will be converted as needed.
model: loaded YOLO model (from load_detector).
conf: minimum confidence threshold.
Returns:
(x1, y1, x2, y2) pixel coordinates of the best detection,
or None if nothing found.
"""
results = model(image, conf=conf, verbose=False)
boxes = results[0].boxes
if boxes is None or len(boxes) == 0:
return None
# Pick highest-confidence box
best = boxes[boxes.conf.argmax()]
x1, y1, x2, y2 = map(int, best.xyxy[0].tolist())
return x1, y1, x2, y2
def crop_roi(image: np.ndarray, box: tuple[int, int, int, int],
padding: float = 0.05) -> np.ndarray:
"""
Crop the ROI from the image with optional relative padding.
Args:
image: BGR or grayscale image.
box: (x1, y1, x2, y2)
padding: fraction of box size to add as margin on each side.
Returns:
Cropped numpy array (same channels as input).
"""
h, w = image.shape[:2]
x1, y1, x2, y2 = box
pw = int((x2 - x1) * padding)
ph = int((y2 - y1) * padding)
x1 = max(0, x1 - pw)
y1 = max(0, y1 - ph)
x2 = min(w, x2 + pw)
y2 = min(h, y2 + ph)
return image[y1:y2, x1:x2]
|