diff --git a/CHANGELOG.md b/CHANGELOG.md index faa2c948..296c0ff7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,11 @@ before 1.4.0 are documented in the ### Added +- **`LibreRFDETRm-ui.pt`: class-agnostic UI element detector (#896).** + UI-DETR-1 (racineai, MIT), an RF-DETR-M fine-tune for screenshots, hosted + as an RF-DETR dataset variant with one class, `object`. Auto-downloads from + `LibreYOLO/LibreRFDETRm-ui`. + - **`val(visualize=True)` draws every validated image for error analysis (#887).** Detection and segmentation images show true positives (green), false positives (red) and false negatives (orange), matched class-aware at @@ -148,6 +153,35 @@ before 1.4.0 are documented in the ### Fixed +- **`Results.plot()` draws boxes, masks, OBB, keypoints and every other + predict output (#896).** It raised `NotImplementedError` for detection, + segmentation, pose, OBB, classification, points, OCR, semantic and panoptic + results. It now returns the annotated image as an `HxWx3` uint8 BGR array, + pixel-identical to `predict(save=True)` (both use the new + `drawing.draw_results`). Accepts `conf`, `labels`, `boxes`, `masks`, + `probs`, `line_width`, `pil`, `img`, `show`, `save` and `filename`. + Classification plots and `predict(save=True)` write the top-5 labels. + Tracked results show their track IDs. Exported-model backends (ONNX and the + other runtimes) save and plot through the same renderer. Dense maps, 3D + cuboids and action chunks keep returning a PIL image unless `pil=False`. + Predict keeps in-memory and URL inputs on the result as `Results.orig_img` + (BGR) so they plot without a path; local files are not kept and `plot()` + reopens them, so directory predictions stay light. A finite video or GIF + collected without `stream=True` keeps no frames; `plot()` decodes the + result's frame on demand. + +- **RF-DETR predict and val resize without antialiasing (#896).** The PIL + bilinear resize antialiased on downscale, unlike RF-DETR training (cv2 + bilinear) and upstream `predict()` since rf-detr 1.9.0. Boxes drifted on + inputs much larger than the canvas (median IoU 0.71 to 0.88 against upstream + on UI screenshots); they now match upstream on screenshots and COCO images, + in PyTorch and ONNX. COCO mAP50-95 of RF-DETR-M on a 200-image val subset + moves from 0.6195 to 0.6179. + +- **Single-class upstream RF-DETR checkpoints convert as `nc=1` (#896).** A + one-output class head (a one-category training set) was converted as + `nc=80` with 79 placeholder names. Predictions were unaffected. + - **Classification `val()`, INT8 calibration and exports use the model's own eval pipeline (#886).** The validator now takes the transform from the model (`_get_eval_transform`), the classification counterpart of diff --git a/docs/nomenclature.md b/docs/nomenclature.md index eedcfe60..ba2e972e 100644 --- a/docs/nomenclature.md +++ b/docs/nomenclature.md @@ -533,7 +533,7 @@ Detector-factory family support follows: | `rtmdet` | `("detect", "segment")` (default: detect) | detect | RTMDet-Ins uses `-seg`; detect training is implemented and directly callable, segment training is not implemented | | `picodet` | `("detect",)` (default) | detect | detect-only | | `ppyoloe` | `("detect",)` (default) | detect | detect-only; no task suffix; pretrained weights linked from the source CDN, not mirrored | -| `rfdetr` | `("detect", "segment", "pose", "obb")` | detect | seg uses smaller sizes; pose/OBB use detect sizes | +| `rfdetr` | `("detect", "segment", "pose", "obb")` | detect | seg uses smaller sizes; pose/OBB use detect sizes; dataset-variant weights `-ui` (class-agnostic UI elements, nc=1) | | `dinov2` | `("semantic", "classify", "embed")` | semantic | DINOv2 backbone + task head; embed bypasses heads and returns the 384-d final CLS token at 224 (all sizes share DINOv2-S); no text tower | | `eomt` | `("semantic", "segment", "panoptic")` | semantic | DINOv2 backbone; sizes s/b/l. Semantic: ADE20K 150-class at 512. Instance segment: COCO 80-class at 640 (l also at 1280). Panoptic: COCO 133-class at 640. Upstream ships no COCO instance checkpoint at s/b. DINOv3 variants excluded | | `pidnet` | `("semantic",)` | semantic | real-time PIDNet semantic segmentation; s/m/l at 1024; Cityscapes 19-class checkpoints; inference + `val`; not trainable in LibreYOLO | @@ -649,6 +649,7 @@ LibreRFDETRn.pt # detect LibreRFDETRn-seg.pt # segment LibreRFDETRx-pose.pt # pose (preview; only size x ships) LibreRFDETRn-obb.pt # obb +LibreRFDETRm-ui.pt # detect; UI element variant (UI-DETR-1, 1 class) # dinov2 — DINOv2 backbone + task head (NOT the RF-DETR detector) LibreDINOv2n.pt # semantic (default task; dense head at 518) diff --git a/libreyolo/backends/base.py b/libreyolo/backends/base.py index a1f96c01..cf78a2fb 100644 --- a/libreyolo/backends/base.py +++ b/libreyolo/backends/base.py @@ -44,14 +44,7 @@ ) from ..preprocess.yolox import preprocess_image as yolox_preprocess_image from ..tasks import normalize_supported_tasks, normalize_task, resolve_task -from ..utils.drawing import ( - draw_boxes, - draw_keypoints, - draw_masks, - draw_obb, - draw_points, - draw_semantic_mask, -) +from ..utils.drawing import draw_results from ..utils.general import ( COCO_CLASSES, get_safe_stem, @@ -62,6 +55,7 @@ from ..utils.model_info import build_model_info, format_model_info from ..utils.predict_args import normalize_predict_kwargs from ..utils.results import ( + keep_source, Boxes, DepthMap, EdgeMap, @@ -4091,80 +4085,12 @@ def _build_result( def _save_annotated(self, result, original_img, image_path, output_path): """Save annotated image to disk.""" - annotated_img = original_img forced_ext = None - if result.boxes is None and getattr(result, "probs", None) is not None: - pass - elif result.boxes is None and getattr(result, "restored", None) is not None: - annotated_img = Image.fromarray(result.restored.array, mode="RGB") - elif result.boxes is None and getattr(result, "depth_map", None) is not None: - from ..utils.drawing import draw_depth_map - - depth_data = result.depth_map.data - if isinstance(depth_data, torch.Tensor): - depth_data = depth_data.cpu().numpy() - annotated_img = draw_depth_map(original_img, depth_data) - elif result.boxes is None and getattr(result, "normal_map", None) is not None: - from ..utils.drawing import draw_normal_map - - normal_data = result.normal_map.data - if isinstance(normal_data, torch.Tensor): - normal_data = normal_data.cpu().numpy() - annotated_img = draw_normal_map(original_img, normal_data) - elif result.boxes is None and getattr(result, "edges", None) is not None: - from ..utils.drawing import draw_edge_map - - edge_data = result.edges.data - if isinstance(edge_data, torch.Tensor): - edge_data = edge_data.cpu().numpy() - annotated_img = draw_edge_map(original_img, edge_data) - elif ( - result.boxes is None and getattr(result, "semantic_mask", None) is not None - ): - mask_data = result.semantic_mask.data - if isinstance(mask_data, torch.Tensor): - mask_data = mask_data.cpu().numpy() - annotated_img = draw_semantic_mask(original_img, mask_data) - elif result.boxes is None and getattr(result, "matte", None) is not None: + if result.boxes is None and getattr(result, "matte", None) is not None: annotated_img = Image.fromarray(result.cutout(original_img), mode="RGBA") forced_ext = "png" - elif result.boxes is None and getattr(result, "points", None) is not None: - if len(result.points) > 0: - annotated_img = draw_points( - original_img, - result.points.xy.tolist(), - result.points.conf.tolist(), - result.points.cls.tolist(), - class_names=result.names, - ) - elif len(result) > 0: - if result.masks is not None: - annotated_img = draw_masks( - annotated_img, - result.masks.data.numpy(), - result.boxes.cls.tolist(), - ) - if result.obb is not None: - annotated_img = draw_obb( - annotated_img, - result.obb.xywhr.tolist(), - result.obb.conf.tolist(), - result.obb.cls.tolist(), - class_names=self.names, - ) - else: - annotated_img = draw_boxes( - annotated_img, - result.boxes.xyxy.tolist(), - result.boxes.conf.tolist(), - result.boxes.cls.tolist(), - class_names=self.names, - ) - if result.keypoints is not None: - kpts_np = result.keypoints.data - if isinstance(kpts_np, torch.Tensor): - kpts_np = kpts_np.cpu().numpy() - annotated_img = draw_keypoints(annotated_img, kpts_np) + else: + annotated_img = draw_results(result, original_img) ext = forced_ext or ( "png" if isinstance(getattr(self, "input_profile", None), dict) @@ -4558,6 +4484,11 @@ def val( # Inference pipeline # ========================================================================= + @staticmethod + def _keep_source(result, original_img, source=None): + """Keep the decoded source image on a result so ``plot()`` works.""" + return keep_source(result, original_img, source) + def _predict_single( self, image: Union[str, Path, Image.Image, np.ndarray], @@ -4602,12 +4533,16 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) if self.task == "embed": - return self._build_embedding_result( - all_outputs, - orig_shape=orig_shape, - image_path=image_path, + return self._keep_source( + self._build_embedding_result( + all_outputs, + orig_shape=orig_shape, + image_path=image_path, + ), + original_img, + image_path, ) if self.task == "restore": result = self._build_restore_result( @@ -4623,7 +4558,7 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) if self.task == "depth": result = self._build_depth_result( all_outputs, @@ -4638,7 +4573,7 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) if self.task == "normal": result = self._build_normal_result( all_outputs, @@ -4653,7 +4588,7 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) if self.task == "edge": result = self._build_edge_result( all_outputs, @@ -4668,7 +4603,7 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) if self.task == "matte": result = self._build_matte_result( all_outputs, @@ -4683,7 +4618,7 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) if self.task == "gaze": result = self._build_gaze_result( all_outputs, @@ -4697,7 +4632,7 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) if self.task == "semantic": result = self._build_semantic_result( all_outputs, @@ -4714,7 +4649,7 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) if self.task == "point": result = self._build_point_result( all_outputs, @@ -4732,7 +4667,7 @@ def _predict_single( image_path if image_path is not None else save_stem, output_path, ) - return result + return self._keep_source(result, original_img, image_path) parsed = self._parse_outputs( all_outputs, @@ -4769,7 +4704,7 @@ def _predict_single( output_path, ) - return result + return self._keep_source(result, original_img, image_path) def _supports_batched_inference(self) -> bool: """Whether ``_run_inference`` accepts stacked (N, C, H, W) blobs. @@ -5063,7 +4998,7 @@ def _predict_batch( if save: self._save_annotated(result, original_img, save_name, output_path) - results.append(result) + results.append(self._keep_source(result, original_img, image_path)) return results # ========================================================================= @@ -5243,7 +5178,7 @@ def _predict_video( source_label = str(source) if source_label is None else source_label effective_imgsz = self._resolve_predict_imgsz(imgsz) - def predict_frame(pil_img): + def predict_frame_result(pil_img): input_tensor, original_img, original_size, ratio = self._preprocess( pil_img, effective_imgsz, "rgb" ) @@ -5349,6 +5284,9 @@ def predict_frame(pil_img): max_det=max_det, ) + def predict_frame(pil_img): + return self._keep_source(predict_frame_result(pil_img), pil_img) + yield from run_video_inference( source, predict_frame, diff --git a/libreyolo/backends/tensorrt.py b/libreyolo/backends/tensorrt.py index f150e2c3..410a941a 100644 --- a/libreyolo/backends/tensorrt.py +++ b/libreyolo/backends/tensorrt.py @@ -522,7 +522,7 @@ def _process_in_batches( if save: self._save_annotated(result, orig_img, save_name, output_path) - results.append(result) + results.append(self._keep_source(result, orig_img, image_path)) return results diff --git a/libreyolo/models/autoconvert.py b/libreyolo/models/autoconvert.py index e045042d..781e32d0 100644 --- a/libreyolo/models/autoconvert.py +++ b/libreyolo/models/autoconvert.py @@ -645,6 +645,14 @@ def _rfdetr_class_metadata( # COCO arch-classes (91 outputs incl. background) -> LibreYOLO's COCO-80. return 80, _checkpoint_names(loaded, 80) + if raw_nc == 0: + # RF-DETR scores classes with independent sigmoids (focal loss); there + # is no background logit. The usual extra output is an unused index + # slot for COCO-style ids, but upstream sizes the head from the + # dataset's category count, so a one-category dataset yields a single + # output and logit 0 is that class (upstream predicts class_id 0). + return 1, _checkpoint_names(loaded, 1) + nc = raw_nc if raw_nc else 80 return nc, _checkpoint_names(loaded, nc) diff --git a/libreyolo/models/base/inference.py b/libreyolo/models/base/inference.py index d0b3fc3f..39a2772d 100644 --- a/libreyolo/models/base/inference.py +++ b/libreyolo/models/base/inference.py @@ -24,22 +24,13 @@ import numpy as np import torch +from PIL import Image from torchvision.ops import batched_nms from ...postprocess.slicing import slice_batch_outputs from ...utils.drawing import ( draw_boxes, - draw_keypoints, - draw_masks, - draw_obb, - draw_depth_map, - draw_edge_map, - draw_normal_map, - draw_mesh, - draw_ocr_regions, - draw_panoptic, - draw_points, - draw_semantic_mask, + draw_results, draw_tile_grid, ) from ...utils.general import ( @@ -51,6 +42,7 @@ from ...utils.image_loader import ImageInput, ImageLoader from ...utils.predict_args import normalize_predict_kwargs from ...utils.results import ( + keep_source, AlbedoMap, Boxes, DepthMap, @@ -592,9 +584,14 @@ def __call__( ) if save: image_path = source if isinstance(source, (str, Path)) else None - img_pil = ImageLoader.load(source, color_format=color_format) ext = output_file_format or "jpg" save_path = resolve_save_path(output_path, image_path, ext=ext) + # Reuse the decoded source rather than fetching the input again. + img_pil = ( + Image.fromarray(result.orig_img[..., ::-1]) + if result.orig_img is not None + else ImageLoader.load(source, color_format=color_format) + ) self._save_annotated_image(result, img_pil, save_path) return result @@ -771,7 +768,11 @@ def _process_in_batches( image if save_stem is None else save_stem, ext=ext, ) - img_pil = ImageLoader.load(image, color_format=color_format) + img_pil = ( + Image.fromarray(result.orig_img[..., ::-1]) + if result.orig_img is not None + else ImageLoader.load(image, color_format=color_format) + ) self._save_annotated_image(result, img_pil, save_path) results.append(result) else: @@ -885,6 +886,7 @@ def _predict_batch( ) image_path = image if isinstance(image, (str, Path)) else None result = self._wrap_results(detections, original_size, image_path, classes) + keep_source(result, original_img, image_path) if save: ext = output_file_format or "jpg" save_path = resolve_save_path( @@ -955,6 +957,8 @@ def _predict_video_clip( ) result = self._wrap_results(detections, original_size, str(source), classes) result.path = str(source) + # A whole-clip result plots on its first sampled frame. + result.orig_img = frames[0] if save: save_file = resolve_save_path( @@ -966,63 +970,17 @@ def _predict_video_clip( def _save_annotated_image( self, result: Results, original_img, save_path: Path ) -> None: - """Internal helper to draw boxes, masks, and keypoints and save to disk.""" - # Classification and whole-image embed results carry no boxes; there is - # nothing to draw, so persist the source image as-is. - if result.boxes is None and ( - getattr(result, "probs", None) is not None - or getattr(result, "embeddings", None) is not None + """Internal helper to render a result on its image and save to disk.""" + # Whole-image embed results carry nothing to draw; persist the source + # image as-is. Classification draws its top-5 in draw_results. + if ( + result.boxes is None + and result.probs is None + and getattr(result, "embeddings", None) is not None ): original_img.save(save_path) log_saved_result(result, save_path) return - if result.boxes is None and getattr(result, "semantic_mask", None) is not None: - mask_data = result.semantic_mask.data - if isinstance(mask_data, torch.Tensor): - mask_data = mask_data.cpu().numpy() - annotated_img = draw_semantic_mask(original_img, mask_data) - annotated_img.save(save_path) - log_saved_result(result, save_path) - return - if result.boxes is None and getattr(result, "panoptic", None) is not None: - pan_data = result.panoptic.data - if isinstance(pan_data, torch.Tensor): - pan_data = pan_data.cpu().numpy() - annotated_img = draw_panoptic( - original_img, - pan_data, - result.panoptic.segments_info, - class_names=result.names, - ) - annotated_img.save(save_path) - log_saved_result(result, save_path) - return - if result.boxes is None and getattr(result, "depth_map", None) is not None: - depth_data = result.depth_map.data - if isinstance(depth_data, torch.Tensor): - depth_data = depth_data.cpu().numpy() - if not result.depth_map.near_is_high: - depth_data = -depth_data - annotated_img = draw_depth_map(original_img, depth_data) - annotated_img.save(save_path) - log_saved_result(result, save_path) - return - if result.boxes is None and getattr(result, "edges", None) is not None: - edge_data = result.edges.data - if isinstance(edge_data, torch.Tensor): - edge_data = edge_data.cpu().numpy() - annotated_img = draw_edge_map(original_img, edge_data) - annotated_img.save(save_path) - log_saved_result(result, save_path) - return - if result.boxes is None and getattr(result, "normal_map", None) is not None: - normal_data = result.normal_map.data - if isinstance(normal_data, torch.Tensor): - normal_data = normal_data.cpu().numpy() - annotated_img = draw_normal_map(original_img, normal_data) - annotated_img.save(save_path) - log_saved_result(result, save_path) - return if result.boxes is None and getattr(result, "albedo", None) is not None: result.albedo.save(save_path) log_saved_result(result, save_path) @@ -1039,87 +997,8 @@ def _save_annotated_image( result.save(png_path, image=original_img) log_saved_result(result, png_path) return - if result.boxes is None and getattr(result, "ocr", None) is not None: - if len(result.ocr) > 0: - ocr_np = result.ocr.numpy() - annotated_img = draw_ocr_regions( - original_img, - ocr_np.data, - ocr_np.texts, - ocr_np.conf, - ) - else: - annotated_img = original_img.copy() - annotated_img.save(save_path) - log_saved_result(result, save_path) - return - if result.boxes is None and getattr(result, "points", None) is not None: - if len(result.points) > 0: - annotated_img = draw_points( - original_img, - result.points.xy.tolist(), - result.points.conf.tolist(), - result.points.cls.tolist(), - class_names=result.names, - ) - else: - annotated_img = original_img.copy() - annotated_img.save(save_path) - log_saved_result(result, save_path) - return - if len(result) > 0: - annotated_img = original_img - # Draw masks first (underneath boxes) - if result.masks is not None: - masks_np = result.masks.data - if isinstance(masks_np, torch.Tensor): - masks_np = masks_np.cpu().numpy() - annotated_img = draw_masks( - annotated_img, - masks_np, - result.boxes.cls.tolist(), - ) - # Draw boxes - if result.obb is not None: - annotated_img = draw_obb( - annotated_img, - result.obb.xywhr.tolist(), - result.obb.conf.tolist(), - result.obb.cls.tolist(), - class_names=result.names, - track_ids=result.obb.id.tolist() - if result.obb.id is not None - else None, - ) - else: - annotated_img = draw_boxes( - annotated_img, - result.boxes.xyxy.tolist(), - result.boxes.conf.tolist(), - result.boxes.cls.tolist(), - class_names=result.names, - ) - # Draw keypoints - if result.keypoints is not None: - kpts_np = result.keypoints.data - if isinstance(kpts_np, torch.Tensor): - kpts_np = kpts_np.cpu().numpy() - annotated_img = draw_keypoints(annotated_img, kpts_np) - # Draw body meshes: projected vertices plus the skeleton through - # the projected joints. - if result.meshes is not None and len(result.meshes) > 0: - meshes_np = result.meshes.numpy() - annotated_img = draw_mesh( - annotated_img, - joints2d=meshes_np.joints2d, - vertices2d=meshes_np.extras.get("vertices2d"), - faces=meshes_np.faces, - vertices3d=meshes_np.vertices, - ) - else: - annotated_img = original_img.copy() - annotated_img.save(save_path) + draw_results(result, original_img).save(save_path) log_saved_result(result, save_path) @staticmethod @@ -1570,6 +1449,7 @@ def _predict_single( # Wrap into Results result = self._wrap_results(detections, original_size, image_path, classes) + keep_source(result, original_img, image_path) # Save annotated image if save: @@ -1658,7 +1538,11 @@ def predict_frame(pil_img): classes=classes, **kwargs, ) - return self._wrap_results(detections, original_size, source_label, classes) + result = self._wrap_results( + detections, original_size, source_label, classes + ) + result.orig_img = original_img + return result return predict_frame @@ -1911,6 +1795,7 @@ def _predict_tiled( "num_detections": len(final_boxes), } result = self._wrap_results(detections, original_size, image_path, classes) + keep_source(result, img_pil, image_path) # Attach tiling metadata as extra attributes result.tiled = True diff --git a/libreyolo/models/base/model.py b/libreyolo/models/base/model.py index 0773ddb3..6232ece1 100644 --- a/libreyolo/models/base/model.py +++ b/libreyolo/models/base/model.py @@ -1298,6 +1298,7 @@ def _predict_augment( from PIL import Image as PILImage from ...utils.image_loader import ImageLoader + from ...utils.results import keep_source effective_imgsz = imgsz if imgsz is not None else self._get_input_size() img_pil = ImageLoader.load(image, color_format=color_format) @@ -1305,7 +1306,7 @@ def _predict_augment( orig_w, orig_h = img_pil.size if getattr(self, "task", "detect") == "semantic": - return self._predict_augment_semantic( + result = self._predict_augment_semantic( img_pil, image_path, (orig_w, orig_h), @@ -1313,9 +1314,10 @@ def _predict_augment( color_format, **kwargs, ) + return keep_source(result, img_pil, image_path) if getattr(self, "task", "detect") == "panoptic": - return self._predict_augment_panoptic( + result = self._predict_augment_panoptic( img_pil, image_path, (orig_w, orig_h), @@ -1323,6 +1325,7 @@ def _predict_augment( color_format, **kwargs, ) + return keep_source(result, img_pil, image_path) scales = (1.0,) if self.TTA_FIXED_SIZE else self.TTA_SCALES @@ -1352,9 +1355,13 @@ def _predict_augment( aug_dets.append((det, orig_size, is_flipped, scale)) if getattr(self, "task", "detect") == "classify": - return self._merge_classify_tta(aug_dets, image_path, (orig_w, orig_h)) - - return self._merge_tta(aug_dets, iou, image_path, (orig_w, orig_h), classes) + result = self._merge_classify_tta(aug_dets, image_path, (orig_w, orig_h)) + else: + result = self._merge_tta( + aug_dets, iou, image_path, (orig_w, orig_h), classes + ) + # Keep the decoded source so plot()/save never fetch the input again. + return keep_source(result, img_pil, image_path) def _postprocess_semantic_logits( self, diff --git a/libreyolo/models/l2cs/inference.py b/libreyolo/models/l2cs/inference.py index 1dfff51e..0255032e 100644 --- a/libreyolo/models/l2cs/inference.py +++ b/libreyolo/models/l2cs/inference.py @@ -18,7 +18,7 @@ from ...utils.general import log_saved_result, resolve_save_path from ...utils.image_loader import ImageInput, ImageLoader -from ...utils.results import Boxes, Gaze, Results +from ...utils.results import Boxes, Gaze, Results, keep_source from ...utils.video import collect_video_results, is_video_file, run_video_inference from .face import FaceBox, FaceDetector, default_face_detector, resolve_face_detector from .utils import bin_logits_to_angles, crop_face, preprocess_face_crops @@ -152,6 +152,7 @@ def _predict_single( faces = self._collect_faces(rgb_np, detector, face_boxes, face_conf) result = self._run_gaze(pil, rgb_np, faces, orig_shape, image_path) + keep_source(result, pil, image) if save: ext = (output_file_format or "jpg").lower().lstrip(".") @@ -176,7 +177,9 @@ def predict_frame(pil_img: Image.Image) -> Results: rgb_np = np.asarray(pil_img) h, w = rgb_np.shape[:2] faces = self._collect_faces(rgb_np, detector, None, face_conf) - return self._run_gaze(pil_img, rgb_np, faces, (h, w), str(source)) + result = self._run_gaze(pil_img, rgb_np, faces, (h, w), str(source)) + result.orig_img = pil_img + return result def annotate(pil_img: Image.Image, result: Results) -> Image.Image: return self._annotate(pil_img, result) @@ -317,29 +320,9 @@ def _empty() -> Results: # ========================================================================= def _annotate(self, pil_img: Image.Image, result: Results) -> Image.Image: - if result.boxes is None or len(result.boxes) == 0: - return pil_img - from ...utils.drawing import draw_boxes, draw_gaze_arrows - - boxes_xyxy = result.boxes.xyxy.tolist() - confs = result.boxes.conf.tolist() - clses = result.boxes.cls.tolist() - annotated = draw_boxes( - pil_img, - boxes_xyxy, - confs, - clses, - class_names=result.names, - ) - if result.gaze is not None and len(result.gaze) > 0: - gaze_np = result.gaze.numpy() if isinstance(result.gaze.data, torch.Tensor) else result.gaze - annotated = draw_gaze_arrows( - annotated, - boxes_xyxy, - gaze_np.data[:, 0], - gaze_np.data[:, 1], - ) - return annotated + from ...utils.drawing import draw_results + + return draw_results(result, pil_img) def _save_annotated_image( self, diff --git a/libreyolo/models/ppocr/inference.py b/libreyolo/models/ppocr/inference.py index b8ced664..a12a1c61 100644 --- a/libreyolo/models/ppocr/inference.py +++ b/libreyolo/models/ppocr/inference.py @@ -20,7 +20,7 @@ from ...postprocess.ppocr import ctc_decode, db_postprocess, sort_quads_reading_order from ...utils.general import log_saved_result, resolve_save_path from ...utils.image_loader import ImageInput, ImageLoader -from ...utils.results import OCRRegions, Results +from ...utils.results import OCRRegions, Results, keep_source from ...utils.video import collect_video_results, is_video_file, run_video_inference from .preprocess import det_normalize, det_resize, get_rotate_crop_image, rec_batches @@ -170,6 +170,7 @@ def _predict_single( rec_batch=rec_batch, image_path=image_path, ) + keep_source(result, pil, image) if save: ext = (output_file_format or "jpg").lower().lstrip(".") save_path = resolve_save_path( @@ -193,13 +194,15 @@ def _predict_video( output_path: Optional[str], ) -> Generator[Results, None, None]: def predict_frame(pil_img: Image.Image) -> Results: - return self.run_pipeline( + result = self.run_pipeline( pil_img, conf=conf, imgsz=imgsz, rec_batch=rec_batch, image_path=str(source), ) + result.orig_img = pil_img + return result yield from run_video_inference( source, diff --git a/libreyolo/models/rfdetr/model.py b/libreyolo/models/rfdetr/model.py index 1f94349e..2d6d61a3 100644 --- a/libreyolo/models/rfdetr/model.py +++ b/libreyolo/models/rfdetr/model.py @@ -14,7 +14,7 @@ from ..base import BaseModel from ...data import load_data_config from ...data.pose_metadata import keypoints_per_class -from ...tasks import normalize_task +from ...tasks import normalize_task, task_to_suffix from ...utils.image_loader import ImageInput, ImageLoader from ...utils.serialization import load_trusted_torch_file from .nn import ( @@ -103,6 +103,14 @@ class LibreRFDETR(BaseModel): # Class-level metadata FAMILY: ClassVar[str] = "rfdetr" FILENAME_PREFIX: ClassVar[str] = "LibreRFDETR" + # Dataset-variant weights: ``-ui`` is the class-agnostic UI element + # detector (UI-DETR-1, racineai, MIT), an RF-DETR-M fine-tune. + WEIGHT_VARIANTS: ClassVar[tuple[str, ...]] = ("ui",) + # The filename regex parses any size/task/variant combination; only these + # (size, task, variant) artifacts are published. + PUBLISHED_WEIGHT_VARIANTS: ClassVar[frozenset[tuple[str, str, str]]] = frozenset( + {("m", "detect", "ui")} + ) # Forward is pure tensor work with no host sync, verified to capture and # replay bit-identically (tests/unit/test_cuda_graph_families.py). SUPPORTS_CUDA_GRAPH = True @@ -331,6 +339,20 @@ def get_download_url(cls, filename: str) -> Optional[str]: upstream_url = _RFDETR_UPSTREAM_WEIGHT_URLS.get(Path(filename).name.lower()) if upstream_url is not None: return upstream_url + variant = cls.detect_variant_from_filename(filename) + if variant is not None: + size = cls.detect_size_from_filename(filename) + task = cls.detect_task_from_filename(filename) or "detect" + if (size, task, variant) not in cls.PUBLISHED_WEIGHT_VARIANTS: + published = ", ".join( + f"{cls.FILENAME_PREFIX}{s}" + f"{'-' + task_to_suffix(t) if task_to_suffix(t) else ''}-{v}.pt" + for s, t, v in sorted(cls.PUBLISHED_WEIGHT_VARIANTS) + ) + raise FileNotFoundError( + f"No published RF-DETR weights for {Path(filename).name}. " + f"Published dataset variants: {published}." + ) return super().get_download_url(filename) # ========================================================================= diff --git a/libreyolo/preprocess/rfdetr.py b/libreyolo/preprocess/rfdetr.py index cb6d6539..b717c6e1 100644 --- a/libreyolo/preprocess/rfdetr.py +++ b/libreyolo/preprocess/rfdetr.py @@ -10,7 +10,6 @@ from typing import Tuple import numpy as np -from PIL import Image IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] @@ -21,7 +20,10 @@ def preprocess_numpy( """ Preprocess RGB HWC uint8 image for RF-DETR inference. - Simple resize + ImageNet normalization. + Simple resize + ImageNet normalization. The resize is bilinear without + antialiasing, matching the cv2 resize used in training and upstream + RF-DETR's ``predict()`` since v1.9.0 (roboflow/rf-detr#1206). Antialiasing + drifts boxes and scores on inputs much larger than the model canvas. Args: img_rgb_hwc: Input image as RGB HWC uint8 numpy array. @@ -30,10 +32,10 @@ def preprocess_numpy( Returns: Tuple of (preprocessed CHW float32 array with ImageNet norm, ratio). """ - img_resized = Image.fromarray(img_rgb_hwc).resize( - (input_size, input_size), Image.Resampling.BILINEAR - ) - arr = np.array(img_resized, dtype=np.float32) / 255.0 + import cv2 + + arr = np.asarray(img_rgb_hwc, dtype=np.float32) / 255.0 + arr = cv2.resize(arr, (input_size, input_size), interpolation=cv2.INTER_LINEAR) mean = np.array(IMAGENET_MEAN, dtype=np.float32) std = np.array(IMAGENET_STD, dtype=np.float32) arr = (arr - mean) / std diff --git a/libreyolo/utils/drawing.py b/libreyolo/utils/drawing.py index e8a30dbc..86624018 100644 --- a/libreyolo/utils/drawing.py +++ b/libreyolo/utils/drawing.py @@ -52,6 +52,29 @@ def get_class_color(class_id: int) -> str: return f"#{r:02x}{g:02x}{b:02x}" +def _box_label( + class_names: list[str] | dict[int, str] | None, + cls_id: int, + score: float, + labels: bool = True, + conf: bool = True, +) -> str | None: + """Return the ``"name: 0.87"`` box label, or the requested part of it.""" + if not labels and not conf: + return None + score_text = f"{float(score):.2f}" + if not labels: + return score_text + class_name = None + if isinstance(class_names, dict): + class_name = class_names.get(cls_id) + elif class_names and cls_id < len(class_names): + class_name = class_names[cls_id] + if class_name is None: + class_name = f"Class {cls_id}" + return f"{class_name}: {score_text}" if conf else str(class_name) + + def draw_boxes( img: Image.Image, boxes: List, @@ -59,6 +82,9 @@ def draw_boxes( classes: List, class_names: List[str] | Dict[int, str] | None = None, track_ids: List | None = None, + labels: bool = True, + conf: bool = True, + line_width: int | None = None, ) -> Image.Image: """ Draw bounding boxes on image with class-specific colors. @@ -75,6 +101,9 @@ def draw_boxes( ID or a dict mapping class ID to class name (default: COCO_CLASSES) track_ids: Optional list of track IDs. When provided, each box is colored by its track ID and the label includes ``ID:``. + labels: Draw class names in box labels. + conf: Draw confidence scores in box labels. + line_width: Box line width in pixels; scales with the image if None. Returns: Annotated PIL Image @@ -89,7 +118,9 @@ def draw_boxes( img_width, img_height = img.size max_dim = max(img_width, img_height) scale_factor = max_dim / 640.0 - box_thickness = max(2, int(2 * scale_factor)) + box_thickness = ( + int(line_width) if line_width else max(2, int(2 * scale_factor)) + ) font_size = max(12, int(12 * scale_factor)) font = _get_font(font_size) @@ -110,12 +141,17 @@ def draw_boxes( ) draw.rectangle([x1, y1, x2, y2], outline=color, width=box_thickness) + if not labels and not conf: + continue # Tracking mode: short two-tone label "#23 0.87" # Detection mode: full label "person: 0.87" if tid is not None: - id_text = f"#{int(tid)}" - conf_text = f" {score:.2f}" + # The track ID is the label in tracking mode. + id_text = f"#{int(tid)}" if labels else "" + conf_text = f"{score:.2f}" if conf else "" + if id_text and conf_text: + conf_text = " " + conf_text # Measure both parts separately for two-tone rendering. id_bbox = draw.textbbox((0, 0), id_text, font=font) full_label = id_text + conf_text @@ -124,16 +160,7 @@ def draw_boxes( text_height = full_bbox[3] - full_bbox[1] id_width = id_bbox[2] - id_bbox[0] else: - class_name = None - if isinstance(class_names, dict): - class_name = class_names.get(cls_id_int) - elif class_names and cls_id_int < len(class_names): - class_name = class_names[cls_id_int] - - if class_name is not None: - full_label = f"{class_name}: {score:.2f}" - else: - full_label = f"Class {cls_id_int}: {score:.2f}" + full_label = _box_label(class_names, cls_id_int, score, labels, conf) full_bbox = draw.textbbox((0, 0), full_label, font=font) text_width = full_bbox[2] - full_bbox[0] text_height = full_bbox[3] - full_bbox[1] @@ -209,6 +236,9 @@ def draw_obb( classes: Sequence[float], class_names: List[str] | Dict[int, str] | None = None, track_ids: Sequence[float] | None = None, + labels: bool = True, + conf: bool = True, + line_width: int | None = None, ) -> Image.Image: """Draw oriented bounding boxes as rotated polygons.""" img_draw = img.copy() @@ -220,7 +250,9 @@ def draw_obb( img_width, img_height = img.size max_dim = max(img_width, img_height) scale_factor = max_dim / 640.0 - box_thickness = max(2, int(2 * scale_factor)) + box_thickness = ( + int(line_width) if line_width else max(2, int(2 * scale_factor)) + ) font_size = max(12, int(12 * scale_factor)) label_padding = max(2, int(2 * scale_factor)) font = _get_font(font_size) @@ -236,20 +268,16 @@ def draw_obb( points = _xywhr_to_points(row) draw.line(points + [points[0]], fill=color, width=box_thickness) + if not labels and not conf: + continue if tid is not None: - label = f"#{int(tid)} {float(score):.2f}" - else: - class_name = None - if isinstance(class_names, dict): - class_name = class_names.get(cls_id_int) - elif class_names and cls_id_int < len(class_names): - class_name = class_names[cls_id_int] - label = ( - f"{class_name}: {float(score):.2f}" - if class_name is not None - else f"Class {cls_id_int}: {float(score):.2f}" + parts = ([f"#{int(tid)}"] if labels else []) + ( + [f"{float(score):.2f}"] if conf else [] ) + label = " ".join(parts) + else: + label = _box_label(class_names, cls_id_int, score, labels, conf) full_bbox = draw.textbbox((0, 0), label, font=font) text_width = full_bbox[2] - full_bbox[0] @@ -1301,3 +1329,158 @@ def draw_actions( if label_y + font.size < bottom: draw.text((8, label_y), names[d], fill=color, font=font) return canvas + + +def _to_numpy(value): + """Tensor or array-like to a numpy array without importing torch.""" + if hasattr(value, "detach"): + value = value.detach().cpu().numpy() + return np.asarray(value) + + +def draw_probs( + img: Image.Image, + top5: Sequence[int], + top5conf: Sequence[float], + class_names: list[str] | dict[int, str] | None = None, + conf: bool = True, +) -> Image.Image: + """Write the top-5 classification labels in the top-left corner.""" + img_draw = img.copy() + draw = ImageDraw.Draw(img_draw) + scale_factor = max(img.size) / 640.0 + font = _get_font(max(12, int(12 * scale_factor))) + padding = max(2, int(2 * scale_factor)) + y = padding + for cls_id, score in zip(top5, top5conf): + cls_id = int(cls_id) + name = _box_label(class_names, cls_id, score, labels=True, conf=False) + text = f"{name} {float(score):.2f}" if conf else name + box = draw.textbbox((0, 0), text, font=font) + width, height = box[2] - box[0], box[3] - box[1] + draw.rectangle( + [0, y - padding, width + padding * 2, y + height + padding * 2], + fill=(0, 0, 0), + ) + draw.text((padding, y), text, fill="white", font=font) + y += height + padding * 3 + return img_draw + + +def draw_results( + result, + img: Image.Image, + *, + conf: bool = True, + labels: bool = True, + boxes: bool = True, + masks: bool = True, + probs: bool = True, + line_width: int | None = None, +) -> Image.Image: + """Render a ``Results`` object on its source image. + + The one renderer behind ``Results.plot()`` and ``predict(save=True)``, so a + plotted result matches the saved file. ``img`` is the RGB source image. + """ + if result.boxes is None and getattr(result, "semantic_mask", None) is not None: + return draw_semantic_mask(img, _to_numpy(result.semantic_mask.data)) + if result.boxes is None and getattr(result, "panoptic", None) is not None: + return draw_panoptic( + img, + _to_numpy(result.panoptic.data), + result.panoptic.segments_info, + class_names=result.names, + ) + if result.boxes is None and getattr(result, "depth_map", None) is not None: + depth = _to_numpy(result.depth_map.data) + if not result.depth_map.near_is_high: + depth = -depth + return draw_depth_map(img, depth) + if result.boxes is None and getattr(result, "edges", None) is not None: + return draw_edge_map(img, _to_numpy(result.edges.data)) + if result.boxes is None and getattr(result, "normal_map", None) is not None: + return draw_normal_map(img, _to_numpy(result.normal_map.data)) + if result.boxes is None and getattr(result, "albedo", None) is not None: + return Image.fromarray(result.albedo.to_rgb()) + if result.boxes is None and getattr(result, "restored", None) is not None: + return Image.fromarray(result.restored.array) + if result.boxes is None and getattr(result, "matte", None) is not None: + return draw_matte(img, result.matte.array) + if result.boxes is None and getattr(result, "ocr", None) is not None: + if len(result.ocr) == 0: + return img.copy() + ocr = result.ocr.numpy() + return draw_ocr_regions(img, ocr.data, ocr.texts, ocr.conf) + if result.boxes is None and getattr(result, "points", None) is not None: + if len(result.points) == 0: + return img.copy() + return draw_points( + img, + result.points.xy.tolist(), + result.points.conf.tolist(), + result.points.cls.tolist(), + class_names=result.names, + ) + if result.boxes is None: + if probs and getattr(result, "probs", None) is not None: + return draw_probs( + img, + result.probs.top5, + _to_numpy(result.probs.top5conf).tolist(), + class_names=result.names, + conf=conf, + ) + return img.copy() + if len(result) == 0: + return img.copy() + + annotated = img + # Masks go underneath boxes. + if masks and result.masks is not None: + annotated = draw_masks( + annotated, _to_numpy(result.masks.data), result.boxes.cls.tolist() + ) + if boxes and result.obb is not None: + annotated = draw_obb( + annotated, + result.obb.xywhr.tolist(), + result.obb.conf.tolist(), + result.obb.cls.tolist(), + class_names=result.names, + track_ids=result.obb.id.tolist() if result.obb.id is not None else None, + labels=labels, + conf=conf, + line_width=line_width, + ) + elif boxes: + annotated = draw_boxes( + annotated, + result.boxes.xyxy.tolist(), + result.boxes.conf.tolist(), + result.boxes.cls.tolist(), + class_names=result.names, + track_ids=result.boxes.id.tolist() if result.boxes.id is not None else None, + labels=labels, + conf=conf, + line_width=line_width, + ) + if result.keypoints is not None: + annotated = draw_keypoints(annotated, _to_numpy(result.keypoints.data)) + if getattr(result, "gaze", None) is not None and len(result.gaze) > 0: + gaze = _to_numpy(result.gaze.data) + annotated = draw_gaze_arrows( + annotated, result.boxes.xyxy.tolist(), gaze[:, 0], gaze[:, 1] + ) + if result.meshes is not None and len(result.meshes) > 0: + meshes = result.meshes.numpy() + annotated = draw_mesh( + annotated, + joints2d=meshes.joints2d, + vertices2d=meshes.extras.get("vertices2d"), + faces=meshes.faces, + vertices3d=meshes.vertices, + ) + if annotated is img: + annotated = img.copy() + return annotated diff --git a/libreyolo/utils/results.py b/libreyolo/utils/results.py index 9b37e242..8335c402 100644 --- a/libreyolo/utils/results.py +++ b/libreyolo/utils/results.py @@ -1898,6 +1898,7 @@ def __init__( boxes3d: Optional[Boxes3D] = None, albedo: Optional[AlbedoMap] = None, actions: Optional[Actions] = None, + orig_img: Any = None, ): if boxes is not None and boxes.orig_shape is None: boxes = boxes.with_orig_shape(orig_shape) @@ -1955,6 +1956,22 @@ def __init__( track_id if track_id is not None else (boxes.id if boxes else None) ) self.frame_idx = frame_idx + self.orig_img = orig_img + + @property + def orig_img(self) -> np.ndarray | None: + """Source image as an ``HxWx3`` uint8 BGR array, or None if not kept.""" + source = self._orig_img + if source is not None and not isinstance(source, np.ndarray): + # Predict keeps the decoded PIL image; convert on first access. + source = np.ascontiguousarray(np.asarray(source.convert("RGB"))[..., ::-1]) + self._orig_img = source + return source + + @orig_img.setter + def orig_img(self, value: Any) -> None: + """Accept a PIL image (RGB) or an ``HxWx3`` BGR array.""" + self._orig_img = value def _new(self, **overrides) -> "Results": data = { @@ -1986,6 +2003,7 @@ def _new(self, **overrides) -> "Results": "speed": dict(self.speed), "track_id": self.track_id, "frame_idx": self.frame_idx, + "orig_img": self._orig_img, } data.update(overrides) return Results(**data) @@ -2143,8 +2161,77 @@ def normals(self) -> Optional[NormalMap]: def normals(self, value: Optional[NormalMap]) -> None: self.normal_map = value - def plot(self, image=None): - """Render dense outputs, calibrated 3D cuboids, or an action chunk.""" + def plot( + self, + image=None, + *, + img=None, + conf: bool = True, + labels: bool = True, + boxes: bool = True, + masks: bool = True, + probs: bool = True, + line_width: int | None = None, + pil: bool | None = None, + show: bool = False, + save: bool = False, + filename: str | None = None, + ): + """Render the result on its source image. + + Boxes, masks, OBB, keypoints, points, OCR, semantic and panoptic maps, + and classification top-5 draw exactly like ``predict(save=True)`` and + return an ``HxWx3`` uint8 BGR array. Dense maps (depth, normal, edge, + albedo), 3D cuboids and action chunks return a PIL image. ``pil`` forces + either return type. + + The source image is ``img`` (BGR array or PIL image), else ``image`` + (RGB array or PIL image), else the image kept by predict + (``orig_img``), else ``path``. + """ + from PIL import Image + + if img is not None: + image = ( + Image.fromarray(np.asarray(img)[..., ::-1].astype(np.uint8)) + if isinstance(img, np.ndarray) + else img + ) + rendered = self._plot_dense(image) + default_pil = rendered is not None + if rendered is None: + from .drawing import draw_results + + source = Image.fromarray(self._source_rgb(image, self.orig_shape)) + rendered = draw_results( + self, + source, + conf=conf, + labels=labels, + boxes=boxes, + masks=masks, + probs=probs, + line_width=line_width, + ) + if show: + rendered.show() + if save or filename: + if filename is None: + source = Path(self.path) if self.path else Path("image.jpg") + if source.suffix.lower() in Image.registered_extensions(): + filename = f"results_{source.name}" + else: + # A video frame or clip: name the still after the video. + frame = f"_{self.frame_idx}" if self.frame_idx is not None else "" + filename = f"results_{source.stem}{frame}.jpg" + Path(filename).parent.mkdir(parents=True, exist_ok=True) + rendered.save(filename) + if pil if pil is not None else default_pil: + return rendered + return np.ascontiguousarray(np.asarray(rendered.convert("RGB"))[..., ::-1]) + + def _plot_dense(self, image=None): + """Render dense outputs, 3D cuboids, or an action chunk; None otherwise.""" if self.actions is not None: from PIL import Image from .drawing import draw_actions @@ -2171,9 +2258,7 @@ def plot(self, image=None): canvas = Image.fromarray(self._source_rgb(image, self.orig_shape)) return draw_boxes3d(canvas, self.boxes3d) if self.normal_map is None and self.edges is None: - raise NotImplementedError( - "Results.plot() supports depth, normal, albedo, edge, and calibrated 3D results." - ) + return None from PIL import Image @@ -2210,13 +2295,24 @@ def _source_rgb(self, image: Any, hw: Tuple[int, int]) -> np.ndarray: from PIL import Image h, w = hw + if image is None and self._orig_img is not None: + image = Image.fromarray(self.orig_img[..., ::-1]) if image is None: + missing = ( + "plot()/cutout()/save() need the source image, but none was kept " + "and Results.path is not an image; pass image= (plot() also takes img=)." + ) if not self.path: - raise ValueError( - "cutout()/save() needs the source image but Results.path is unset; " - "pass image=." - ) - rgb = np.asarray(Image.open(self.path).convert("RGB")) + raise ValueError(missing) + # A frame of a collected video or GIF: path is the whole clip, so + # decode that frame on demand (PIL would open a GIF at frame 0). + rgb = self._video_frame_rgb() + if rgb is None: + try: + rgb = np.asarray(Image.open(self.path).convert("RGB")) + except (OSError, ValueError) as exc: + raise ValueError(missing) from exc elif isinstance(image, Image.Image): rgb = np.asarray(image.convert("RGB")) else: @@ -2231,6 +2327,20 @@ def _source_rgb(self, image: Any, hw: Tuple[int, int]) -> np.ndarray: ) return rgb.astype(np.uint8) + def _video_frame_rgb(self) -> np.ndarray | None: + """Decode frame ``frame_idx`` of the video at ``path``, or None.""" + if self.frame_idx is None or not self.path or not Path(self.path).is_file(): + return None + import cv2 + + cap = cv2.VideoCapture(str(self.path)) + try: + cap.set(cv2.CAP_PROP_POS_FRAMES, int(self.frame_idx)) + ok, frame = cap.read() + finally: + cap.release() + return frame[..., ::-1] if ok else None + def save(self, path: str, image: Any = None) -> str: """Save a matte result as a transparent-background RGBA PNG cutout. @@ -2242,7 +2352,7 @@ def save(self, path: str, image: Any = None) -> str: if self.matte is None: raise NotImplementedError( "Results.save() writes a transparent-PNG cutout and is defined for " - "matte results only. Use result.plot()/CLI --save for other tasks." + "matte results only. Use result.plot(save=True) or CLI --save for other tasks." ) rgba = self.cutout(image=image) out = Path(path) @@ -2690,3 +2800,17 @@ def collect(value: Any) -> None: f"Cannot stack embeddings with different dimensions: {sorted(dimensions)}." ) return torch.cat(non_empty, dim=0) + + +def keep_source(result: Any, image: Any, source: Any = None) -> Any: + """Attach the decoded source image so ``Results.plot()`` can draw on it. + + Skipped when ``source`` is a local image file: ``plot()`` reopens it from + ``Results.path``, so a directory or list prediction does not hold every + decoded image in memory. In-memory inputs and URLs keep their pixels. + """ + if isinstance(result, Results) and not ( + isinstance(source, (str, Path)) and Path(source).is_file() + ): + result.orig_img = image + return result diff --git a/libreyolo/utils/video.py b/libreyolo/utils/video.py index efb9d0dd..b93e4bff 100644 --- a/libreyolo/utils/video.py +++ b/libreyolo/utils/video.py @@ -390,7 +390,14 @@ def collect_video_results( f"Consider using stream=True to avoid high memory usage.", stacklevel=3, ) - return list(gen) + results = [] + for result in gen: + # A collected list must not hold every decoded frame; stream=True + # keeps each frame's source image for plotting. + if getattr(result, "_orig_img", None) is not None: + result.orig_img = None + results.append(result) + return results def run_video_inference( @@ -419,29 +426,18 @@ def run_video_inference( show: Display frames in a cv2 window. output_path: Output path for saved video. annotate_fn: Optional callable ``(pil_img, result) -> pil_img`` for - custom annotation (e.g. tracking labels). When *None*, the default - ``draw_boxes()`` annotation is used. + custom annotation (e.g. tracking labels). When *None*, the frame is + drawn with ``draw_results()``, as ``Results.plot()`` does. progress: Show a tqdm progress bar (frames processed, fps). Yields: ``Results`` for each processed frame. """ import cv2 - import torch from PIL import Image from tqdm import tqdm - from .drawing import ( - draw_boxes, - draw_depth_map, - draw_edge_map, - draw_normal_map, - draw_keypoints, - draw_masks, - draw_matte, - draw_obb, - draw_points, - ) + from .drawing import draw_results if isinstance(source, (str, Path)): frame_source = VideoSource(source, vid_stride=vid_stride) @@ -528,107 +524,8 @@ def run_video_inference( if save or show: if annotate_fn is not None: annotated_pil = annotate_fn(pil_img, result) - elif ( - result.boxes is None - and getattr(result, "probs", None) is not None - ): - annotated_pil = pil_img - elif ( - result.boxes is None - and getattr(result, "points", None) is not None - ): - if len(result.points) > 0: - annotated_pil = draw_points( - pil_img, - result.points.xy.tolist(), - result.points.conf.tolist(), - result.points.cls.tolist(), - class_names=result.names, - ) - else: - annotated_pil = pil_img - elif ( - result.boxes is None - and getattr(result, "restored", None) is not None - ): - annotated_pil = Image.fromarray( - result.restored.array, mode="RGB" - ) - elif result.boxes is None and getattr(result, "albedo", None) is not None: - annotated_pil = Image.fromarray(result.albedo.to_rgb()) - elif ( - result.boxes is None - and getattr(result, "matte", None) is not None - ): - # Checkerboard-composited cutout preview (video frames - # cannot carry an alpha channel, so the transparency is - # visualized instead). - annotated_pil = draw_matte(pil_img, result.matte.array) - elif ( - result.boxes is None - and getattr(result, "depth_map", None) is not None - ): - depth_np = result.depth_map.data - if isinstance(depth_np, torch.Tensor): - depth_np = depth_np.cpu().numpy() - if not result.depth_map.near_is_high: - depth_np = -depth_np - annotated_pil = draw_depth_map(pil_img, depth_np) - elif ( - result.boxes is None - and getattr(result, "normal_map", None) is not None - ): - normal_np = result.normal_map.data - if isinstance(normal_np, torch.Tensor): - normal_np = normal_np.cpu().numpy() - annotated_pil = draw_normal_map(pil_img, normal_np) - elif ( - result.boxes is None - and getattr(result, "edges", None) is not None - ): - edge_np = result.edges.data - if isinstance(edge_np, torch.Tensor): - edge_np = edge_np.cpu().numpy() - annotated_pil = draw_edge_map(pil_img, edge_np) - elif len(result) > 0: - annotated_pil = pil_img - if result.masks is not None: - masks_np = result.masks.data - if isinstance(masks_np, torch.Tensor): - masks_np = masks_np.cpu().numpy() - annotated_pil = draw_masks( - annotated_pil, - masks_np, - result.boxes.cls.tolist(), - ) - if result.obb is not None: - annotated_pil = draw_obb( - annotated_pil, - result.obb.xywhr.tolist(), - result.obb.conf.tolist(), - result.obb.cls.tolist(), - class_names=result.names, - track_ids=( - result.obb.id.tolist() - if result.obb.id is not None - else None - ), - ) - else: - annotated_pil = draw_boxes( - annotated_pil, - result.boxes.xyxy.tolist(), - result.boxes.conf.tolist(), - result.boxes.cls.tolist(), - class_names=result.names, - ) - if result.keypoints is not None: - kpts_np = result.keypoints.data - if isinstance(kpts_np, torch.Tensor): - kpts_np = kpts_np.cpu().numpy() - annotated_pil = draw_keypoints(annotated_pil, kpts_np) else: - annotated_pil = pil_img + annotated_pil = draw_results(result, pil_img) annotated_bgr = cv2.cvtColor( np.array(annotated_pil), cv2.COLOR_RGB2BGR diff --git a/skills/libreyolo-upload-hf-model/SKILL.md b/skills/libreyolo-upload-hf-model/SKILL.md index 5d9745a7..172a540b 100644 --- a/skills/libreyolo-upload-hf-model/SKILL.md +++ b/skills/libreyolo-upload-hf-model/SKILL.md @@ -214,7 +214,7 @@ LibreRFDETRn.pt, LibreRFDETRs.pt, LibreRFDETRm.pt, LibreRFDETRl.pt, LibreRFDETRn-seg.pt, LibreRFDETRs-seg.pt, LibreRFDETRm-seg.pt, LibreRFDETRl-seg.pt, LibreRFDETRx-pose.pt, LibreRFDETRn-obb.pt, LibreRFDETRs-obb.pt, LibreRFDETRm-obb.pt, -LibreRFDETRl-obb.pt, +LibreRFDETRl-obb.pt, LibreRFDETRm-ui.pt, LibreDETRr50.pt, LibreDETRr50dc5.pt, LibreDETRr101.pt, LibreDETRr101dc5.pt, diff --git a/tests/unit/test_autoconvert.py b/tests/unit/test_autoconvert.py index 78d237f3..7c1716c1 100644 --- a/tests/unit/test_autoconvert.py +++ b/tests/unit/test_autoconvert.py @@ -376,6 +376,14 @@ def test_rfdetr_custom_90_class_names_are_not_coerced_to_coco(self): 90, ) == (90, names) + def test_rfdetr_single_logit_head_is_one_class(self): + # A one-category upstream dataset yields a single-output head with no + # spare slot (raw_nc 0); it must not fall back to COCO's 80 classes. + assert autoconvert_module._rfdetr_class_metadata( + {"args": argparse.Namespace(num_classes=1, class_names=["object"])}, + 0, + ) == (1, ["object"]) + def test_rfdetr_coco_metadata_maps_90_arch_classes_to_coco80(self): assert autoconvert_module._rfdetr_class_metadata( {"args": argparse.Namespace(dataset_file="coco")}, diff --git a/tests/unit/test_l2cs.py b/tests/unit/test_l2cs.py index b2a39f9e..f82e57ca 100644 --- a/tests/unit/test_l2cs.py +++ b/tests/unit/test_l2cs.py @@ -188,6 +188,22 @@ def test_libre_l2cs_end_to_end_with_byo_bbox(tmp_path): assert torch.isfinite(result.gaze.data).all() +def test_libre_l2cs_plot_in_memory_source(tmp_path): + """plot() works on a gaze result from an in-memory image (no path).""" + weights_path = tmp_path / "LibreL2CSr18.pt" + torch.save(_make_dummy_state_dict("r18"), weights_path) + model = LibreL2CS(str(weights_path), size="r18", device="cpu") + + img = np.full((48, 64, 3), 128, dtype=np.uint8) + result = model(img, face_boxes=[(0, 0, 64, 48)]) + + assert result.path is None + assert result.orig_img.shape == (48, 64, 3) + plotted = result.plot() + assert isinstance(plotted, np.ndarray) + assert plotted.shape == (48, 64, 3) + + def test_libre_l2cs_callable_face_detector(tmp_path): """A user-supplied callable becomes the face detector.""" sd = _make_dummy_state_dict("r18") diff --git a/tests/unit/test_ppocr.py b/tests/unit/test_ppocr.py index 0cb3226c..c8f2c011 100644 --- a/tests/unit/test_ppocr.py +++ b/tests/unit/test_ppocr.py @@ -306,6 +306,22 @@ def test_predict_without_charset_raises(): model.predict(img) +def test_predict_in_memory_source_can_plot(): + """plot() works on an OCR result from an in-memory image (no path).""" + from libreyolo.models.ppocr.model import PPOCR_V5_NUM_CLASSES, LibrePPOCR + + model = LibrePPOCR(model_path=None, size="t") + model.charset = ["x"] * PPOCR_V5_NUM_CLASSES + img = np.full((48, 64, 3), 255, dtype=np.uint8) + result = model.predict(img) + + assert result.path is None + assert result.orig_img.shape == (48, 64, 3) + plotted = result.plot() + assert isinstance(plotted, np.ndarray) + assert plotted.shape == (48, 64, 3) + + # -------------------------------------------------------------------------- # Dataset resolver on the committed fixture # -------------------------------------------------------------------------- diff --git a/tests/unit/test_results.py b/tests/unit/test_results.py index 94fac453..7727b5a6 100644 --- a/tests/unit/test_results.py +++ b/tests/unit/test_results.py @@ -722,3 +722,199 @@ def test_draw_semantic_mask_paints_classes_and_skips_ignore(): assert pixels[0, 0].any() # class-1 half painted assert not pixels[0, 9].any() # ignore half untouched (still black) + + +def _detect_result(orig_img=None, path=None): + return Results( + boxes=Boxes( + torch.tensor([[8.0, 10.0, 40.0, 44.0], [30.0, 5.0, 60.0, 30.0]]), + torch.tensor([0.9, 0.6]), + torch.tensor([0.0, 2.0]), + ), + orig_shape=(48, 64), + path=path, + names={0: "person", 1: "bicycle", 2: "car"}, + orig_img=orig_img, + ) + + +def _source_rgb(): + return np.random.default_rng(0).integers(0, 256, (48, 64, 3), dtype=np.uint8) + + +class TestResultsPlot: + def test_detect_plot_returns_bgr_array_matching_saved_render(self, tmp_path): + from libreyolo.utils.drawing import draw_results + + rgb = _source_rgb() + result = _detect_result(orig_img=Image.fromarray(rgb)) + + plotted = result.plot() + + expected = np.asarray(draw_results(result, Image.fromarray(rgb)))[..., ::-1] + assert plotted.dtype == np.uint8 and plotted.shape == (48, 64, 3) + np.testing.assert_array_equal(plotted, expected) + assert not np.array_equal(plotted, rgb[..., ::-1]) # boxes were drawn + + # predict(save=True) goes through the same renderer. + saved = tmp_path / "saved.png" + InferenceRunner._save_annotated_image(None, result, Image.fromarray(rgb), saved) + np.testing.assert_array_equal( + np.asarray(Image.open(saved).convert("RGB"))[..., ::-1], plotted + ) + + def test_plot_pil_and_source_image_options(self): + rgb = _source_rgb() + result = _detect_result(orig_img=Image.fromarray(rgb)) + default = result.plot() + + as_pil = result.plot(pil=True) + assert isinstance(as_pil, Image.Image) + np.testing.assert_array_equal(np.asarray(as_pil)[..., ::-1], default) + + bare = _detect_result() + np.testing.assert_array_equal(bare.plot(img=rgb[..., ::-1].copy()), default) + np.testing.assert_array_equal(bare.plot(Image.fromarray(rgb)), default) + + def test_plot_toggles_and_save(self, tmp_path): + rgb = _source_rgb() + result = _detect_result(orig_img=Image.fromarray(rgb)) + default = result.plot() + + assert not np.array_equal(result.plot(conf=False, labels=False), default) + np.testing.assert_array_equal(result.plot(boxes=False), rgb[..., ::-1]) + + out = tmp_path / "sub" / "plot.png" + result.plot(filename=str(out)) + assert out.exists() + + def test_plot_without_source_image_says_how_to_fix(self): + with pytest.raises(ValueError, match="img="): + _detect_result().plot() + + def test_orig_img_is_bgr_and_survives_slicing(self): + rgb = _source_rgb() + result = _detect_result(orig_img=Image.fromarray(rgb)) + + np.testing.assert_array_equal(result.orig_img, rgb[..., ::-1]) + np.testing.assert_array_equal(result[:1].orig_img, rgb[..., ::-1]) + + def test_classify_plot_writes_top5(self): + rgb = np.full((48, 64, 3), 200, dtype=np.uint8) + result = Results( + boxes=None, + orig_shape=(48, 64), + names={0: "cat", 1: "dog"}, + probs=Probs(torch.tensor([0.2, 0.8])), + orig_img=Image.fromarray(rgb), + ) + + assert not np.array_equal(result.plot(), rgb) + np.testing.assert_array_equal(result.plot(probs=False), rgb) + + +class TestResultsPlotFollowUps: + def test_plot_with_video_path_and_no_frame_says_how_to_fix(self, tmp_path): + clip = tmp_path / "clip.mp4" + clip.write_bytes(b"not an image") + + with pytest.raises(ValueError, match="img="): + _detect_result(path=str(clip)).plot() + + def test_tracked_plot_draws_track_ids(self): + from libreyolo.utils.drawing import draw_boxes + + rgb = _source_rgb() + result = _detect_result(orig_img=Image.fromarray(rgb)) + result.boxes = result.boxes.with_id(torch.tensor([7.0, 9.0])) + + expected = draw_boxes( + Image.fromarray(rgb), + result.boxes.xyxy.tolist(), + result.boxes.conf.tolist(), + result.boxes.cls.tolist(), + class_names=result.names, + track_ids=[7.0, 9.0], + ) + np.testing.assert_array_equal(result.plot(pil=True), np.asarray(expected)) + + def test_tracked_obb_label_options(self): + rgb = Image.fromarray(_source_rgb()) + args = ([[30.0, 24.0, 20.0, 10.0, 0.3]], [0.8], [1.0]) + full = np.asarray(draw_obb(rgb, *args, track_ids=[3])) + no_conf = np.asarray(draw_obb(rgb, *args, track_ids=[3], conf=False)) + no_label = np.asarray(draw_obb(rgb, *args, track_ids=[3], labels=False)) + + assert not np.array_equal(full, no_conf) + assert not np.array_equal(full, no_label) + assert not np.array_equal(no_conf, no_label) + + def test_classify_save_matches_plot(self, tmp_path): + rgb = np.full((48, 64, 3), 200, dtype=np.uint8) + result = Results( + boxes=None, + orig_shape=(48, 64), + names={0: "cat", 1: "dog"}, + probs=Probs(torch.tensor([0.2, 0.8])), + orig_img=Image.fromarray(rgb), + ) + saved = tmp_path / "cls.png" + InferenceRunner._save_annotated_image(None, result, Image.fromarray(rgb), saved) + + np.testing.assert_array_equal( + np.asarray(Image.open(saved).convert("RGB"))[..., ::-1], result.plot() + ) + + +class TestKeepSource: + def test_local_files_are_reopened_not_kept(self, tmp_path): + from libreyolo.utils.results import keep_source + + rgb = _source_rgb() + path = tmp_path / "img.png" + Image.fromarray(rgb).save(path) + + on_disk = keep_source(_detect_result(path=str(path)), Image.fromarray(rgb), str(path)) + in_memory = keep_source(_detect_result(), Image.fromarray(rgb), None) + url = keep_source(_detect_result(), Image.fromarray(rgb), "https://x.test/a.jpg") + + assert on_disk.orig_img is None + assert in_memory.orig_img is not None and url.orig_img is not None + np.testing.assert_array_equal(on_disk.plot(), in_memory.plot()) + + def test_gif_frame_plots_on_its_own_frame(self, tmp_path): + frames = [ + Image.fromarray(np.full((48, 64, 3), 40 * k, dtype=np.uint8)) for k in range(4) + ] + gif = tmp_path / "anim.gif" + frames[0].save(gif, save_all=True, append_images=frames[1:], loop=0) + result = _detect_result(path=str(gif)) + result.frame_idx = 2 + + decoded = np.asarray(result._video_frame_rgb()) + np.testing.assert_array_equal(result.plot(), result.plot(Image.fromarray(decoded))) + assert not np.array_equal(decoded, np.asarray(frames[0])) + + +def test_gaze_plot_draws_arrows_like_l2cs_save(): + from libreyolo.utils.drawing import draw_boxes, draw_gaze_arrows + from libreyolo.utils.results import Gaze + + rgb = _source_rgb() + result = _detect_result(orig_img=Image.fromarray(rgb)) + result.gaze = Gaze(torch.tensor([[0.2, -0.3], [-0.1, 0.4]])) + + boxes = result.boxes.xyxy.tolist() + expected = draw_gaze_arrows( + draw_boxes( + Image.fromarray(rgb), + boxes, + result.boxes.conf.tolist(), + result.boxes.cls.tolist(), + class_names=result.names, + ), + boxes, + np.array([0.2, -0.1], dtype=np.float32), + np.array([-0.3, 0.4], dtype=np.float32), + ) + np.testing.assert_array_equal(result.plot(pil=True), np.asarray(expected)) diff --git a/tests/unit/test_rfdetr_preprocess.py b/tests/unit/test_rfdetr_preprocess.py index 0f820ee1..4f999fb1 100644 --- a/tests/unit/test_rfdetr_preprocess.py +++ b/tests/unit/test_rfdetr_preprocess.py @@ -61,3 +61,61 @@ def test_rfdetr_val_preprocessor_requests_original_image(): preprocessor = RFDETRValPreprocessor((64, 64), max_labels=4) assert preprocessor.wants_unresized_image is True + + +def test_rfdetr_preprocess_resize_is_not_antialiased(): + # Training resizes with cv2 bilinear and upstream predict() uses + # antialias=False; an antialiased resize drifts boxes on large inputs. + import torch.nn.functional as F + + from libreyolo.models.rfdetr.utils import ( + IMAGENET_MEAN, + IMAGENET_STD, + preprocess_numpy, + ) + + rgb = np.random.default_rng(2).integers(0, 256, (211, 419, 3), dtype=np.uint8) + chw, _ = preprocess_numpy(rgb, 64) + + src = torch.from_numpy(rgb.astype(np.float32) / 255.0).permute(2, 0, 1)[None] + ref = F.interpolate( + src, size=(64, 64), mode="bilinear", align_corners=False, antialias=False + )[0] + mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1) + std = torch.tensor(IMAGENET_STD).view(3, 1, 1) + torch.testing.assert_close( + torch.from_numpy(chw), (ref - mean) / std, atol=1e-3, rtol=0 + ) + + +def test_rfdetr_ui_weight_variant_filename(): + from libreyolo.models.rfdetr.model import LibreRFDETR + + name = "LibreRFDETRm-ui.pt" + assert LibreRFDETR.detect_size_from_filename(name) == "m" + assert LibreRFDETR.detect_task_from_filename(name) is None + assert LibreRFDETR.detect_variant_from_filename(name) == "ui" + assert LibreRFDETR.get_download_url(name) == ( + "https://huggingface.co/LibreYOLO/LibreRFDETRm-ui/resolve/main/" + "LibreRFDETRm-ui.pt" + ) + + +@pytest.mark.parametrize( + "name", ["LibreRFDETRs-ui.pt", "LibreRFDETRm-seg-ui.pt", "LibreRFDETRl-obb-ui.pt"] +) +def test_rfdetr_unpublished_ui_variant_names_are_rejected(name): + from libreyolo.models.rfdetr.model import LibreRFDETR + + with pytest.raises(FileNotFoundError, match="LibreRFDETRm-ui.pt"): + LibreRFDETR.get_download_url(name) + + +def test_rfdetr_non_variant_download_urls_unchanged(): + from libreyolo.models.rfdetr.model import LibreRFDETR + + assert LibreRFDETR.get_download_url("LibreRFDETRs-seg.pt") == ( + "https://huggingface.co/LibreYOLO/LibreRFDETRs-seg/resolve/main/" + "LibreRFDETRs-seg.pt" + ) + assert LibreRFDETR.get_download_url("LibreYOLO9s.pt") is None diff --git a/tests/unit/test_video.py b/tests/unit/test_video.py index 7e42f381..c5d3ab01 100644 --- a/tests/unit/test_video.py +++ b/tests/unit/test_video.py @@ -448,3 +448,38 @@ def test_warns_on_large_video(self, tmp_path): collect_video_results(iter([1, 2, 3]), path, vid_stride=1) memory_warnings = [x for x in w if "stream=True" in str(x.message)] assert len(memory_warnings) == 1 + + +def test_collected_video_results_drop_source_frames(sample_video): + from PIL import Image + + from libreyolo.utils.results import Results + + frames = [ + Results(boxes=None, orig_shape=(4, 4), orig_img=Image.new("RGB", (4, 4))) + for _ in range(3) + ] + collected = collect_video_results(iter(frames), sample_video, vid_stride=1) + + assert [r.orig_img for r in collected] == [None, None, None] + + +def test_collected_video_frame_plots_by_decoding_its_frame(sample_video, tmp_path, monkeypatch): + import torch + + from libreyolo.utils.results import Boxes, Results + + frames = {idx: frame for frame, idx in VideoSource(sample_video)} + result = Results( + boxes=Boxes(torch.tensor([[4.0, 4.0, 30.0, 30.0]]), torch.tensor([0.9]), torch.tensor([0.0])), + orig_shape=(64, 64), + path=sample_video, + names={0: "thing"}, + frame_idx=4, + ) + + np.testing.assert_array_equal(result.plot(), result.plot(img=frames[4])) + + monkeypatch.chdir(tmp_path) + result.plot(save=True) + assert (tmp_path / "results_test_video_4.jpg").exists() diff --git a/weights/LICENSE_NOTICE.txt b/weights/LICENSE_NOTICE.txt index c8f09434..1d939259 100644 --- a/weights/LICENSE_NOTICE.txt +++ b/weights/LICENSE_NOTICE.txt @@ -97,6 +97,12 @@ project the weights were derived from. Per-family summary: upstream: roboflow/rf-detr (backbone: facebookresearch/dinov2, Apache-2.0) + RF-DETR UI (LibreRFDETRm-ui) MIT + upstream: racineai/UI-DETR-1 (HF revision 0f0dda5), an + RF-DETR-M fine-tune on Apache-2.0 base weights. Training data: + six Roboflow Universe UI datasets, declared MIT by the authors. + Public mirror: https://huggingface.co/LibreYOLO/LibreRFDETRm-ui + Faster R-CNN (LibreFasterRCNN{n,s,m,l}) BSD-3-Clause (implied) upstream code: pytorch/vision v0.26.0 at commit 336d36e8db990a905498c73933e35231876e28bc (BSD-3-Clause)