from pathlib import Path from shutil import rmtree import cv2 import hydra import numpy as np import torch from omegaconf import DictConfig from torchvision.ops import box_iou from tqdm import tqdm from src.dl.dataset import Loader, parse_yolo_label_file from src.dl.utils import ( abs_xyxy_to_norm_xywh, get_latest_experiment_name, norm_xywh_to_abs_xyxy, vis_one_box, ) from src.infer.torch_model import Torch_model def norm_xywh_to_xyxy(norm_boxes: torch.Tensor) -> torch.Tensor: # boxes: [N,4] in (cx, cy, w, h), all in [0,1] cx, cy, w, h = norm_boxes.unbind(-1) x1 = cx - w / 2 y1 = cy - h / 2 x2 = cx + w / 2 y2 = cy + h / 2 return torch.stack([x1, y1, x2, y2], dim=-1).clamp(0, 1) def save_case( case_type: str, img: np.ndarray, abs_box: torch.Tensor, # [4] absolute xyxy label: int, score: float | None, img_path: Path, output_dir: Path, label_to_name: dict, ): # Draw and save image + single-line YOLO txt with the corresponding box out_img_dir = output_dir / case_type out_img_dir.mkdir(parents=True, exist_ok=True) mode = "pred" if case_type == "FPs" else "gt" draw = img.copy() vis_one_box( draw, abs_box, int(label), mode=mode, label_to_name=label_to_name, score=score if score is not None else None, ) cv2.imwrite(str(out_img_dir / f"{Path(img_path).stem}.jpg"), draw) def check_results( img, img_path, preds, targets, iou_thresh: float, conf_thresh: float, output_dir: Path, label_to_name: dict, ): """ Match predictions to targets: - same class - IoU >= iou_thresh (computed on normalized xyxy) - pred score >= conf_thresh Save unmatched predictions as FP and unmatched targets as FN. """ H, W = img.shape[:2] # Prepare tensors if len(preds) == 0: pred_boxes_xyxy_abs = torch.zeros((0, 4), dtype=torch.float32) pred_boxes_norm_xywh = torch.zeros((0, 4), dtype=torch.float32) pred_labels = torch.zeros((0,), dtype=torch.long) pred_scores = torch.zeros((0,), dtype=torch.float32) else: pred_boxes_xyxy_abs = torch.as_tensor(preds["boxes"], dtype=torch.float32) # abs xyxy pred_boxes_norm_xywh = torch.as_tensor( preds["norm_boxes"], dtype=torch.float32 ) # norm xywh pred_labels = torch.as_tensor(preds["labels"], dtype=torch.long) pred_scores = torch.as_tensor(preds["scores"], dtype=torch.float32) if len(targets["boxes"]) == 0: tgt_boxes_norm_xywh = torch.zeros((0, 4), dtype=torch.float32) tgt_labels = torch.zeros((0,), dtype=torch.long) else: tgt_boxes_norm_xywh = torch.as_tensor(targets["boxes"], dtype=torch.float32) # norm xywh tgt_labels = torch.as_tensor(targets["labels"], dtype=torch.long) # filter by conf_thresh keep = pred_scores >= conf_thresh pred_boxes_xyxy_abs = pred_boxes_xyxy_abs[keep] pred_boxes_norm_xywh = pred_boxes_norm_xywh[keep] pred_labels = pred_labels[keep] pred_scores = pred_scores[keep] # Early outs if pred_boxes_norm_xywh.numel() == 0 and tgt_boxes_norm_xywh.numel() == 0: return if tgt_boxes_norm_xywh.numel() == 0: # All kept preds become FP for i in range(len(pred_labels)): save_case( "FPs", img, pred_boxes_xyxy_abs[i], int(pred_labels[i]), float(pred_scores[i]), img_path, output_dir, label_to_name, ) return if pred_boxes_norm_xywh.numel() == 0: # All targets become FN tgt_abs_xyxy = norm_xywh_to_abs_xyxy(tgt_boxes_norm_xywh, H, W) for j in range(len(tgt_labels)): save_case( "FNs", img, tgt_abs_xyxy[j], int(tgt_labels[j]), 0, img_path, output_dir, label_to_name, ) return # Compute IoU on normalized xyxy (scale doesn't matter) pred_xyxy_norm = norm_xywh_to_xyxy(pred_boxes_norm_xywh) # [Np,4] tgt_xyxy_norm = norm_xywh_to_xyxy(tgt_boxes_norm_xywh) # [Nt,4] ious = box_iou(pred_xyxy_norm, tgt_xyxy_norm) # [Np,Nt] # Greedy one-to-one matching by pred score Np, Nt = ious.shape tgt_matched = torch.zeros(Nt, dtype=torch.bool) pred_matched = torch.zeros(Np, dtype=torch.bool) order = torch.argsort(pred_scores, descending=True) for pi in order.tolist(): # mask by class and unmatched targets same_cls = tgt_labels == pred_labels[pi] cand = (~tgt_matched) & same_cls & (ious[pi] >= iou_thresh) if cand.any(): # pick best IoU target among candidates ti = torch.argmax(ious[pi] * cand.float()).item() tgt_matched[ti] = True pred_matched[pi] = True # Save FPs and FNs for pi in torch.where(~pred_matched)[0].tolist(): save_case( "FPs", img, pred_boxes_xyxy_abs[pi], int(pred_labels[pi]), float(pred_scores[pi]), img_path, output_dir, label_to_name, ) tgt_abs_xyxy = norm_xywh_to_abs_xyxy(tgt_boxes_norm_xywh, H, W) for ti in torch.where(~tgt_matched)[0].tolist(): save_case( "FNs", img, tgt_abs_xyxy[ti], int(tgt_labels[ti]), 0, img_path, output_dir, label_to_name, ) def run(model, train_loader, val_loader, cfg: DictConfig) -> None: output_dir = Path(cfg.train.infer_path).parent / "check_errors" if output_dir.exists(): rmtree(output_dir) output_dir.mkdir(parents=True, exist_ok=True) for loader in [train_loader, val_loader]: split_name = loader.dataset.mode for batch in tqdm(loader, desc=f"Processing {split_name}"): _, _, paths = batch img = cv2.imread(Path(cfg.train.data_path) / "images" / paths[0]) label_path = Path(cfg.train.data_path) / "labels" / paths[0].with_suffix(".txt") targets = [{"boxes": [], "labels": []}] if label_path.exists() and label_path.stat().st_size > 1: targets_raw, _ = parse_yolo_label_file(label_path) targets = [ { "boxes": np.array(targets_raw[:, 1:5], dtype=np.float32), "labels": np.array(targets_raw[:, 0], dtype=np.int64), } ] preds = model(img) for pred in preds: # Model returns torch tensors, convert to cpu for processing pred["boxes"] = pred["boxes"].cpu() pred["labels"] = pred["labels"].cpu() pred["scores"] = pred["scores"].cpu() pred["norm_boxes"] = abs_xyxy_to_norm_xywh( pred["boxes"], img.shape[0], img.shape[1] ) check_results( img=img, img_path=paths[0], preds=preds[0], targets=targets[0], iou_thresh=cfg.train.iou_thresh, conf_thresh=cfg.train.conf_thresh, output_dir=output_dir, label_to_name=cfg.train.label_to_name, ) @hydra.main(version_base=None, config_path="../../", config_name="config") def main(cfg: DictConfig) -> None: cfg.exp = get_latest_experiment_name(cfg.exp, cfg.train.path_to_save) model = Torch_model( model_name=cfg.model_name, model_path=Path(cfg.train.path_to_save) / "model.pt", n_outputs=len(cfg.train.label_to_name), input_width=cfg.train.img_size[1], input_height=cfg.train.img_size[0], conf_thresh=cfg.train.conf_thresh, rect=cfg.export.dynamic_input, keep_ratio=cfg.train.keep_ratio, use_nms=True, # to remove duplocated boxes on 1 GT object and not show them as FPs ) base_loader = Loader( root_path=Path(cfg.train.data_path), img_size=tuple(cfg.train.img_size), batch_size=1, num_workers=cfg.train.num_workers, cfg=cfg, debug_img_processing=cfg.train.debug_img_processing, ) train_loader, val_loader, _ = base_loader.build_dataloaders() run(model, train_loader, val_loader, cfg) if __name__ == "__main__": main()