68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
|
|
class FailureReason(str, Enum):
|
|
no_face = "no_face"
|
|
small_face = "small_face"
|
|
align_fail = "align_fail"
|
|
infer_fail = "infer_fail"
|
|
read_fail = "read_fail"
|
|
det_fail = "det_fail"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Detection:
|
|
bbox_xyxy: np.ndarray # float32 shape (4,) in original image coords
|
|
landmarks5: np.ndarray # float32 shape (5,2) in original image coords
|
|
score: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageResult:
|
|
ok: bool
|
|
reason: Optional[FailureReason]
|
|
detail: str
|
|
emb: Optional[np.ndarray] = None # float32 (D,) L2-normalized
|
|
det: Optional[Detection] = None
|
|
|
|
|
|
@dataclass
|
|
class BuildReport:
|
|
total_person_dirs: int = 0
|
|
enrolled_persons: int = 0
|
|
total_images: int = 0
|
|
processed_images: int = 0
|
|
ok_images: int = 0
|
|
failed_images: int = 0
|
|
per_person_used: Dict[str, int] = None
|
|
per_person_failed: Dict[str, List[Tuple[str, FailureReason, str]]] = None
|
|
failure_reasons: Dict[str, int] = None
|
|
skipped_persons: List[str] = None
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.per_person_used is None:
|
|
self.per_person_used = {}
|
|
if self.per_person_failed is None:
|
|
self.per_person_failed = {}
|
|
if self.failure_reasons is None:
|
|
self.failure_reasons = {}
|
|
if self.skipped_persons is None:
|
|
self.skipped_persons = []
|
|
|
|
def add_failure(self, person: str, img_path: str, reason: FailureReason, detail: str) -> None:
|
|
self.processed_images += 1
|
|
self.failed_images += 1
|
|
self.failure_reasons[str(reason)] = self.failure_reasons.get(str(reason), 0) + 1
|
|
self.per_person_failed.setdefault(person, []).append((img_path, reason, detail))
|
|
|
|
def add_success(self, person: str) -> None:
|
|
self.processed_images += 1
|
|
self.ok_images += 1
|
|
self.per_person_used[person] = self.per_person_used.get(person, 0) + 1
|