import json import os from typing import Dict, Iterable, List, Optional, Sequence, Tuple import cv2 import numpy as np BBox = Tuple[int, int, int, int] def list_images(folder: str, extensions: Optional[Sequence[str]] = None) -> List[str]: """Liệt kê các ảnh trong thư mục theo phần mở rộng cho trước.""" if extensions is None: extensions = [".jpg", ".jpeg", ".png", ".bmp", ".webp"] paths: List[str] = [] if not os.path.isdir(folder): return paths for name in os.listdir(folder): _, ext = os.path.splitext(name) if ext.lower() in extensions: paths.append(os.path.join(folder, name)) return sorted(paths) def load_image_bgr(path: str) -> np.ndarray: """Đọc ảnh bằng OpenCV và trả về ảnh dạng BGR.""" image = cv2.imread(path, cv2.IMREAD_COLOR) if image is None: raise ValueError(f"Không đọc được ảnh: {path}") return image def bgr_to_rgb(image_bgr: np.ndarray) -> np.ndarray: """Chuyển ảnh từ BGR sang RGB.""" return cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) def rgb_to_bgr(image_rgb: np.ndarray) -> np.ndarray: """Chuyển ảnh từ RGB sang BGR.""" return cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR) def rotate_image(image: np.ndarray, angle: float) -> np.ndarray: """Xoay ảnh theo góc bất kỳ và giữ nguyên toàn bộ nội dung.""" if angle % 360 == 0: return image (h, w) = image.shape[:2] center = (w / 2.0, h / 2.0) matrix = cv2.getRotationMatrix2D(center, angle, 1.0) cos = abs(matrix[0, 0]) sin = abs(matrix[0, 1]) new_w = int((h * sin) + (w * cos)) new_h = int((h * cos) + (w * sin)) matrix[0, 2] += (new_w / 2) - center[0] matrix[1, 2] += (new_h / 2) - center[1] return cv2.warpAffine(image, matrix, (new_w, new_h), flags=cv2.INTER_LINEAR) def trim_white_border( image_bgr: np.ndarray, threshold: int = 245, min_size: int = 5, ) -> np.ndarray: """Cắt viền trắng dư thừa của template trước khi xử lý.""" gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) mask = gray < threshold if not np.any(mask): return image_bgr ys, xs = np.where(mask) y0, y1 = int(ys.min()), int(ys.max()) x0, x1 = int(xs.min()), int(xs.max()) cropped = image_bgr[y0 : y1 + 1, x0 : x1 + 1] if cropped.shape[0] < min_size or cropped.shape[1] < min_size: return image_bgr return cropped def compute_scale_candidates( image_shape: Tuple[int, int, int], template_shape: Tuple[int, int, int], min_scale: float, max_scale: float, steps: int, ) -> List[float]: """Tính danh sách tỉ lệ scale sao cho template vẫn nằm trong ảnh.""" image_h, image_w = image_shape[:2] template_h, template_w = template_shape[:2] max_allowed = min(image_w / template_w, image_h / template_h) if max_allowed <= 0: return [] scale_max = min(max_scale, max_allowed) scale_min = min_scale if scale_max < scale_min: scale_min = scale_max if steps <= 1: return [max(scale_min, 0.01)] scales = np.linspace(scale_min, scale_max, steps) return [float(s) for s in scales if s > 0.01] def resize_template(template: np.ndarray, scale: float) -> np.ndarray: """Resize template theo tỉ lệ scale.""" h, w = template.shape[:2] new_w = max(int(w * scale), 1) new_h = max(int(h * scale), 1) return cv2.resize(template, (new_w, new_h), interpolation=cv2.INTER_AREA) def nms_boxes(boxes: List[BBox], scores: List[float], iou_threshold: float) -> List[int]: """Áp dụng NMS và trả về chỉ số bbox được giữ lại.""" if not boxes: return [] x1 = np.array([b[0] for b in boxes], dtype=np.float32) y1 = np.array([b[1] for b in boxes], dtype=np.float32) x2 = np.array([b[0] + b[2] for b in boxes], dtype=np.float32) y2 = np.array([b[1] + b[3] for b in boxes], dtype=np.float32) scores_np = np.array(scores, dtype=np.float32) areas = (x2 - x1 + 1) * (y2 - y1 + 1) order = scores_np.argsort()[::-1] keep: List[int] = [] while order.size > 0: i = int(order[0]) keep.append(i) xx1 = np.maximum(x1[i], x1[order[1:]]) yy1 = np.maximum(y1[i], y1[order[1:]]) xx2 = np.minimum(x2[i], x2[order[1:]]) yy2 = np.minimum(y2[i], y2[order[1:]]) w = np.maximum(0.0, xx2 - xx1 + 1) h = np.maximum(0.0, yy2 - yy1 + 1) inter = w * h iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-6) remaining = np.where(iou <= iou_threshold)[0] order = order[remaining + 1] return keep def color_from_name(name: str) -> Tuple[int, int, int]: """Sinh màu ổn định dựa trên tên template.""" seed = abs(hash(name)) % 255 r = (seed * 97) % 255 g = (seed * 57) % 255 b = (seed * 17) % 255 return int(b), int(g), int(r) def draw_results( image_bgr: np.ndarray, results: Iterable[Dict], template_colors: Dict[str, Tuple[int, int, int]], ) -> np.ndarray: """Vẽ bbox và điểm cosine similarity (nếu có) lên ảnh.""" output = image_bgr.copy() for item in results: x, y, w, h = item["bbox"] name = item["template_name"] color = template_colors.get(name, (0, 255, 0)) cv2.rectangle(output, (x, y), (x + w, y + h), color, 2) match_score = item.get("match_score") sim = item.get("cosine_similarity") label_parts = [] if match_score is not None: label_parts.append(f"m:{match_score:.2f}") if sim is not None: label_parts.append(f"c:{sim:.2f}") if label_parts: label = " ".join(label_parts) cv2.putText( output, label, (x, max(y - 6, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA, ) return output def to_json(data: List[Dict]) -> str: """Chuyển danh sách kết quả sang chuỗi JSON.""" return json.dumps(data, ensure_ascii=False, indent=2) def is_mostly_white(image_bgr: np.ndarray, threshold: int = 100, min_foreground_ratio: float = 0.01) -> bool: """Kiem tra vung anh co qua it net ve (gan nhu toan trang).""" gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) foreground = gray < threshold ratio = float(foreground.mean()) return ratio < min_foreground_ratio