/j)ddlmZddlmZddlmZddlZddlmZ ddl m Z m Z ddl mZddlmZmZddlmZmZdd lmZGd d eZdS) ) annotations)Path)AnyN)ClassificationDatasetbuild_dataloader) BaseValidator)LOGGERRANK)ClassifyMetricsConfusionMatrix) plot_imagesceZdZdZd&d'fd Zd(d Zd)d Zd*dZd+dZd,dZ d-dZ d.dZ d,dZ d/dZ d0d!Zd,d"Zd1d$Zd2d%ZxZS)3ClassificationValidatoraA class extending the BaseValidator class for validation based on a classification model. This validator handles the validation process for classification models, including metrics calculation, confusion matrix generation, and visualization of results. Attributes: targets (list[torch.Tensor]): Ground truth class labels. pred (list[torch.Tensor]): Model predictions. metrics (ClassifyMetrics): Object to calculate and store classification metrics. names (dict): Mapping of class indices to class names. nc (int): Number of classes. confusion_matrix (ConfusionMatrix): Matrix to evaluate model performance across classes. Methods: get_desc: Return a formatted string summarizing classification metrics. init_metrics: Initialize confusion matrix, class names, and tracking containers. preprocess: Preprocess input batch by moving data to device. update_metrics: Update running metrics with model predictions and batch targets. finalize_metrics: Finalize metrics including confusion matrix and processing speed. postprocess: Extract the primary prediction from model output. get_stats: Calculate and return a dictionary of metrics. build_dataset: Create a ClassificationDataset instance for validation. get_dataloader: Build and return a data loader for classification validation. print_results: Print evaluation metrics for the classification model. plot_val_samples: Plot validation image samples with their ground truth labels. plot_predictions: Plot images with their predicted class labels. Examples: >>> from ultralytics.models.yolo.classify import ClassificationValidator >>> args = dict(model="yolo26n-cls.pt", data="imagenet10") >>> validator = ClassificationValidator(args=args) >>> validator() Notes: Torchvision classification models can also be passed to the 'model' argument, i.e. model='resnet18'. N _callbacks dict | NonereturnNonect||||d|_d|_d|j_t |_dS)aInitialize ClassificationValidator with dataloader, save directory, and other parameters. Args: dataloader (torch.utils.data.DataLoader, optional): DataLoader to use for validation. save_dir (str | Path, optional): Directory to save results. args (dict, optional): Arguments containing model and validation configuration. _callbacks (dict, optional): Dictionary of callback functions to be called during validation. Nclassify)super__init__targetspredargstaskr metrics)self dataloadersave_dirrr __class__s i/home/longshao/multi-rider-rag/.venv/lib/python3.11/site-packages/ultralytics/models/yolo/classify/val.pyrz ClassificationValidator.__init__8sK XtZ@@@  # &(( strc ddzS)z=Return a formatted string summarizing classification metrics.z %22s%11s%11s)classestop1_acctop5_accrs r!get_descz ClassificationValidator.get_descGs#'JJJr"modeltorch.nn.Modulec|j|_t|j|_g|_g|_t |j|_dS)z^Initialize confusion matrix, class names, and tracking containers for predictions and targets.)namesN)r.lenncrrr confusion_matrix)rr+s r! init_metricsz$ClassificationValidator.init_metricsKsE[ ek""  /ek B B Br"batchdict[str, Any]cX|d|j|jjdk|d<|jjr|dn|d|d<|d|j|jjdk|d<|S)zTPreprocess input batch by moving data to device and converting to appropriate dtype.imgcuda) non_blockingcls)todevicetyperhalffloat)rr3s r! preprocessz"ClassificationValidator.preprocessSsU|t{AQU[A[\\e .2inVuU|((***%,BTBTBVBVe U|t{AQU[A[\\e  r"preds torch.Tensorctt|jd}|j|ddddd|ftj |j |dtj dS)aUpdate running metrics with model predictions and batch targets. Args: preds (torch.Tensor): Model predictions, typically logits or probabilities for each class. batch (dict): Batch data containing images and class labels. Notes: This method appends the top-N predictions (sorted by confidence in descending order) to the prediction list for later evaluation. N is limited to the minimum of 5 and the number of classes. T) descendingNr9) minr/r.rappendargsortr<torchint32cpur)rr@r3n5s r!update_metricsz&ClassificationValidator.update_metricsZsTZ! $ $ qT::111crc6BGG TTXXZZ[[[ E%L--ek::>>@@AAAAAr"c&|j|j|j|jjr,dD])}|j|j||j*|j |j _ |j|j _|j|j _dS)aFinalize metrics including confusion matrix and processing speed. Examples: >>> validator = ClassificationValidator() >>> validator.pred = [torch.tensor([[0, 1, 2]])] # Top-3 predictions for one sample >>> validator.targets = [torch.tensor([0])] # Ground truth class >>> validator.finalize_metrics() >>> print(validator.metrics.confusion_matrix) # Access the confusion matrix Notes: This method processes the accumulated predictions and targets to generate the confusion matrix, optionally plots it, and updates the metrics object with speed information. )TF)r normalizeon_plotN) r1process_cls_predsrrrplotsplotrrPspeedr)rrOs r!finalize_metricsz(ClassificationValidator.finalize_metricsis // 4<HHH 9? n( n n %**DMY`d`l*mmmm!Z  $  (,(= %%%r"7torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor]cNt|ttfr|dn|S)zSExtract the primary prediction from model output if it's in a list or tuple format.r) isinstancelisttuple)rr@s r! postprocessz#ClassificationValidator.postprocesss#%edE];;FuQxxFr"dict[str, float]cd|j|j|j|jjS)zSCalculate and return a dictionary of metrics by processing targets and predictions.)rprocessrr results_dictr)s r! get_statsz!ClassificationValidator.get_statss) T\49555|((r"ctdkrdgtjz}dgtjz}tj|j|dtj|j|dd|D|_d|D|_dStdkr:tj|jddtj|jdddSdS)zGather stats from all GPUs.rN)dstcg|] }|D]}| Sr(r().0rankrs r! z8ClassificationValidator.gather_stats..s%JJJ$TJJTJJJJr"cg|] }|D]}| Sr(r()rdrers r!rfz8ClassificationValidator.gather_stats..s&UUUPTUUWGUUUUr")r distget_world_size gather_objectrr)rgathered_predsgathered_targetss r! gather_statsz$ClassificationValidator.gather_statss 199"Vd&9&;&;;N $v(;(=(==   ty.a @ @ @ @  t|-=1 E E E EJJ.JJJDIUU0@UUUDLLL AXX  ty$A 6 6 6 6  t|Tq 9 9 9 9 9 9Xr"img_pathrcFt||jd|jjS)z7Create a ClassificationDataset instance for validation.F)rootraugmentprefix)rrsplit)rrns r! build_datasetz%ClassificationValidator.build_datasets"$(EZ^ZcZijjjjr" dataset_path Path | str batch_sizeinttorch.utils.data.DataLoadercf||}t|||jjdS)aPBuild and return a data loader for classification validation. Args: dataset_path (str | Path): Path to the dataset directory. batch_size (int): Number of samples per batch. Returns: (torch.utils.data.DataLoader): DataLoader object for the classification validation dataset. )re)rtrrworkers)rrurwdatasets r!get_dataloaderz&ClassificationValidator.get_dataloaders3$$\22TY5FRPPPPr"cddt|jjzz}tj|d|jj|jjfzdS)z6Print evaluation metrics for the classification model.z%22sz%11.3gallN)r/rkeysr infotop1top5)rpfs r! print_resultsz%ClassificationValidator.print_resultssJ hT\%6!7!77 7 B%!2DL4EFFGGGGGr"nictj|djd|d<t||jd|dz |j|jdS)aPlot validation image samples with their ground truth labels. Args: batch (dict[str, Any]): Dictionary containing batch data with 'img' (images) and 'cls' (class labels). ni (int): Batch index used for naming the output file. Examples: >>> validator = ClassificationValidator() >>> batch = {"img": torch.rand(16, 3, 224, 224), "cls": torch.randint(0, 10, (16,))} >>> validator.plot_val_samples(batch, 0) r6r batch_idx val_batchz _labels.jpg)labelsfnamer.rPN)rIarangeshaper rr.rP)rr3rs r!plot_val_samplesz(ClassificationValidator.plot_val_samplessh#\%,*>> validator = ClassificationValidator() >>> batch = {"img": torch.rand(16, 3, 224, 224)} >>> preds = torch.rand(16, 10) # 16 images, 10 classes >>> validator.plot_predictions(batch, preds, 0) r6rrD)dim)r6rr9confrz _pred.jpg)rr.rPN) dictrIrrargmaxamaxr rr.rP)rr3r@r batched_predss r!plot_predictionsz(ClassificationValidator.plot_predictionsse l5<#5a#899 U***Eq)))     -";b";";";;*L       r")NNNN)rrrr)rr#)r+r,rr)r3r4rr4)r@rAr3r4rr)rr)r@rVrrA)rr\)rnr#rr)rurvrwrxrry)r3r4rrxrr)r3r4r@rArrxrr)__name__ __module__ __qualname____doc__rr*r2r?rMrUr[r`rmrtr~rrr __classcell__)r s@r!rrsm##J ) ) ) ) ) ) )KKKKCCCC B B B B>>>>,GGGG)))) : : : :kkkk Q Q Q QHHHH     (        r"r) __future__rpathlibrtypingrrItorch.distributed distributedrhultralytics.datarrultralytics.engine.validatorrultralytics.utilsr r ultralytics.utils.metricsr r ultralytics.utils.plottingr rr(r"r!rs#""""" DDDDDDDD666666********FFFFFFFF222222F F F F F mF F F F F r"