Fi#ddlmZddlmZddlmZddlZddlmZddl m Z m Z m Z ddl mZddlmZmZd ZGd d eZGd d eZy)) annotations)Path)AnyN) YOLODataset)ComposeFormat v8_transforms)DetectionValidator)colorstrops)RTDETRValidatorc<eZdZdZddfd Zdfd ZddZxZS) RTDETRDatasetaReal-Time DEtection and TRacking (RT-DETR) dataset class extending the base YOLODataset class. This specialized dataset class is designed for use with the RT-DETR object detection model and is optimized for real-time detection and tracking tasks. Attributes: augment (bool): Whether to apply data augmentation. rect (bool): Whether to use rectangular training. use_segments (bool): Whether to use segmentation masks. use_keypoints (bool): Whether to use keypoint annotations. imgsz (int): Target image size for training. Methods: load_image: Load one image from dataset index. build_transforms: Build transformation pipeline for the dataset. Examples: Initialize an RT-DETR dataset >>> dataset = RTDETRDataset(img_path="path/to/images", imgsz=640) >>> image, hw0, hw = dataset.load_image(0) N)datac(t||d|i|y)aLInitialize the RTDETRDataset class by inheriting from the YOLODataset class. This constructor sets up a dataset specifically optimized for the RT-DETR (Real-Time DEtection and TRacking) model, building upon the base YOLODataset functionality. Args: *args (Any): Variable length argument list passed to the parent YOLODataset class. data (dict | None): Dictionary containing dataset information. If None, default values will be used. **kwargs (Any): Additional keyword arguments passed to the parent YOLODataset class. rN)super__init__)selfrargskwargs __class__s _/var/www/pod-logistic/pod-ai/venv/lib/python3.12/site-packages/ultralytics/models/rtdetr/val.pyrzRTDETRDataset.__init__)s $4T4V4c&t|||S)aLoad one image from dataset index 'i'. Args: i (int): Index of the image to load. rect_mode (bool, optional): Whether to use rectangular mode for batch inference. Returns: im (np.ndarray): Loaded image as a NumPy array. hw_original (tuple[int, int]): Original image dimensions in (height, width) format. hw_resized (tuple[int, int]): Resized image dimensions in (height, width) format. Examples: Load an image from the dataset >>> dataset = RTDETRDataset(img_path="path/to/images") >>> image, hw0, hw = dataset.load_image(0) )i rect_mode)r load_image)rrrrs rrzRTDETRDataset.load_image6s"w!A!;;rc |jr|jr|js |jnd|_|jr|js |jnd|_|jr|js |jnd|_t ||j |d}n tg}|jtdd|j|jd|j|j|S)zBuild transformation pipeline for the dataset. Args: hyp (dict, optional): Hyperparameters for transformations. Returns: (Compose): Composition of transformation functions. gT)stretchxywh) bbox_format normalize return_maskreturn_keypoint batch_idx mask_ratio mask_overlap)augmentrectmosaicmixupcutmixr imgszrappendr use_segments use_keypointsr& overlap_mask)rhyp transformss rbuild_transformszRTDETRDataset.build_transformsIs <<'+||DII3CJ%)\\$)) CI'+||DII3CJ&tTZZdKJ!J " -- $ 2 2>> --  r)F)N)__name__ __module__ __qualname____doc__rrr4 __classcell__)rs@rrrs,$( 5<&rrc8eZdZdZddZddZ d dZd dZy) r aRTDETRValidator extends the DetectionValidator class to provide validation capabilities specifically tailored for the RT-DETR (Real-Time DETR) object detection model. The class allows building of an RTDETR-specific dataset for validation, applies confidence thresholding for post-processing, and updates evaluation metrics accordingly. Attributes: args (Namespace): Configuration arguments for validation. data (dict): Dataset configuration dictionary. Methods: build_dataset: Build an RTDETR Dataset for validation. postprocess: Apply confidence thresholding to prediction outputs. Examples: Initialize and run RT-DETR validation >>> from ultralytics.models.rtdetr import RTDETRValidator >>> args = dict(model="rtdetr-l.pt", data="coco8.yaml") >>> validator = RTDETRValidator(args=args) >>> validator() Notes: For further details on the attributes and methods, refer to the parent DetectionValidator class. Nc t||jj|d|jd|jjxsdt |d|j  S)aBuild an RTDETR Dataset. Args: img_path (str): Path to the folder containing images. mode (str, optional): `train` mode or `val` mode, users are able to customize different augmentations for each mode. batch (int, optional): Size of batches, this is for `rect`. Returns: (RTDETRDataset): Dataset configured for RT-DETR validation. FNz: ) img_pathr- batch_sizer(r2r)cacheprefixr)rrr-r>r r)rr<modebatchs r build_datasetzRTDETRValidator.build_datasetsT))// ))//)TtfBK(  rc|S)zJReturn predictions unchanged as RT-DETR handles scaling in postprocessing.)rprednpbatchs r scale_predszRTDETRValidator.scale_predss rct|ttfs|dg}|dj\}}}|dj d|dz fd\}}||j j z}tjd|jg|z}t|D]\}} tj| } ||jd\} } tj| | d| dgd} | | jd } | | |j j kD||<|D cgc]} | ddddf| dddf| ddd fd !c} Scc} w) aApply confidence thresholding to prediction outputs. Args: preds (torch.Tensor | list | tuple): Raw predictions from the model. If tensor, should have shape (batch_size, num_predictions, num_classes + 4) where last dimension contains bbox coords and class scores. Returns: (list[dict[str, torch.Tensor]]): List of dictionaries for each image, each containing: - 'bboxes': Tensor of shape (N, 4) with bounding box coordinates - 'conf': Tensor of shape (N,) with confidence scores - 'cls': Tensor of shape (N,) with class indices Nr)dim)r)device).NT) descending)bboxesconfcls) isinstancelisttupleshapesplitrr-torchzerosrM enumerater xywh2xyxymaxcatargsortrQ)rpredsbs_ndrPscoresoutputsrbboxscorerRpredxs r postprocesszRTDETRValidator.postprocesssE %$/DME!HNN ArqBF <$))//!;;vfmm<=B ( 6GAt==&Dr*JE399dE)$4c)nE2ND  67Dediinn45GAJ  6PWW!1QU8Qq!tWQq!tWEWWWs!$Ec t|d}|j}|jr t|n|}|dj }|dddgfxx|dd|j j z zcc<|dddgfxx|dd|j j z zcc<tj|}|d d d dfxx|d d dd fdz zcc<t|j|d j|d jD]i\}}} |jj||j|jt| |D cgc]} t| dc} t|d d ky cc} w)aSerialize YOLO predictions to COCO json format. Args: predn (dict[str, torch.Tensor]): Predictions dictionary containing 'bboxes', 'conf', and 'cls' keys with bounding box coordinates, confidence scores, and class predictions. pbatch (dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'. im_filerP.r ori_shapeNrQrRrO)image_id file_name category_idrerf)rstem isnumericintclonerr-r xyxy2xywhziptolistjdictr.name class_mapround) rrErFpathrsrpboxbscrhs r pred_to_jsonzRTDETRValidator.pred_to_jsonscF9%&yy $ 03t9dHo##% C!QKF;/2TYY__DD C!QKF;/2TYY__DDmmC  ArrE c!QR%j1n$ 3::<v)=)=)?uATATAVW GAq! JJ   (!%#'>>#aΛ4QU1a[4"1a[    5sF)valN)rEdict[str, torch.Tensor]rFdict[str, Any]returnr)r_z7torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor]rzlist[dict[str, torch.Tensor]])rErrFrrNone)r5r6r7r8rBrGrirrDrrr r hs/2 0XLX &XBrr ) __future__rpathlibrtypingrrXultralytics.datarultralytics.data.augmentrrr ultralytics.models.yolo.detectr ultralytics.utilsr r __all__rr rDrrrsE# (CC=+ SKSlp(pr