From d3179daa4d28c83a3febfce32d6c8469bb144694 Mon Sep 17 00:00:00 2001 From: Zhaoyan Fang <52028100+satuoqaq@users.noreply.github.com> Date: Wed, 15 Feb 2023 19:05:00 +0800 Subject: [PATCH] YOLOv7 Assigner visualization (#543) * add yolov7 assigner visual * base yolov5 detector wrirte yolov7 * update readme * add yolov7 assigner visual * base yolov5 detector wrirte yolov7 * update * Update projects/assigner_visualization/README.md Co-authored-by: Nioolek <40284075+Nioolek@users.noreply.github.com> * Update projects/assigner_visualization/README.md Co-authored-by: Nioolek <40284075+Nioolek@users.noreply.github.com> * add note and typehint * update --------- Co-authored-by: Nioolek <40284075+Nioolek@users.noreply.github.com> --- projects/assigner_visualization/README.md | 16 +- .../assigner_visualization.py | 24 +-- ...t_8xb16-300e_coco_assignervisualization.py | 9 + .../dense_heads/__init__.py | 3 +- .../dense_heads/yolov7_head_assigner.py | 159 ++++++++++++++++++ .../detectors/yolo_detector_assigner.py | 5 +- 6 files changed, 200 insertions(+), 16 deletions(-) create mode 100644 projects/assigner_visualization/configs/yolov7_tiny_syncbn_fast_8xb16-300e_coco_assignervisualization.py create mode 100644 projects/assigner_visualization/dense_heads/yolov7_head_assigner.py diff --git a/projects/assigner_visualization/README.md b/projects/assigner_visualization/README.md index 2ae5703a..0bf0d8dc 100644 --- a/projects/assigner_visualization/README.md +++ b/projects/assigner_visualization/README.md @@ -6,22 +6,32 @@ This project is developed for easily showing assigning results. The script allows users to analyze where and how many positive samples each gt is assigned in the image. -Now, the script supports `YOLOv5` and `RTMDet`. +Now, the script supports `YOLOv5`, `YOLOv7` and `RTMDet`. ## Usage ### Command +YOLOv5 assigner visualization command: + ```shell python projects/assigner_visualization/assigner_visualization.py projects/assigner_visualization/configs/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_assignervisualization.py ``` Note: `YOLOv5` does not need to load the trained weights. +YOLOv7 assigner visualization command: + +```shell +python projects/assigner_visualization/assigner_visualization.py projects/assigner_visualization/configs/yolov7_tiny_syncbn_fast_8xb16-300e_coco_assignervisualization.py -c ${checkpont} +``` + +RTMdet assigner visualization command: + ```shell python projects/assigner_visualization/assigner_visualization.py projects/assigner_visualization/configs/rtmdet_s_syncbn_fast_8xb32-300e_coco_assignervisualization.py -c ${checkpont} ``` -${checkpont} is the checkpont file path. Dynamic label assignment is used in `RTMDet`, model weights will affect the positive sample allocation results, so it is recommended to load the trained model weights. +${checkpont} is the checkpont file path. Dynamic label assignment is used in `YOLOv7` and `RTMDet`, model weights will affect the positive sample allocation results, so it is recommended to load the trained model weights. -If you want to know details about label assignment, you can check the [documentation](https://mmyolo.readthedocs.io/zh_CN/latest/algorithm_descriptions/rtmdet_description.html#id5). +If you want to know details about label assignment, you can check the [RTMDet](https://mmyolo.readthedocs.io/zh_CN/latest/algorithm_descriptions/rtmdet_description.html#id5). diff --git a/projects/assigner_visualization/assigner_visualization.py b/projects/assigner_visualization/assigner_visualization.py index 139efe5b..f6c7d4f6 100644 --- a/projects/assigner_visualization/assigner_visualization.py +++ b/projects/assigner_visualization/assigner_visualization.py @@ -11,13 +11,14 @@ import torch from mmengine import ProgressBar from mmengine.config import Config, DictAction from mmengine.dataset import COLLATE_FUNCTIONS -from mmengine.runner import load_checkpoint +from mmengine.runner.checkpoint import load_checkpoint from numpy import random from mmyolo.registry import DATASETS, MODELS from mmyolo.utils import register_all_modules from projects.assigner_visualization.dense_heads import (RTMHeadAssigner, - YOLOv5HeadAssigner) + YOLOv5HeadAssigner, + YOLOv7HeadAssigner) from projects.assigner_visualization.visualization import \ YOLOAssignerVisualizer @@ -87,17 +88,20 @@ def main(): # build model model = MODELS.build(cfg.model) if args.checkpoint is not None: - _ = load_checkpoint(model, args.checkpoint, map_location='cpu') - elif isinstance(model.bbox_head, RTMHeadAssigner): + load_checkpoint(model, args.checkpoint) + elif isinstance(model.bbox_head, (YOLOv7HeadAssigner, RTMHeadAssigner)): warnings.warn( - 'if you use dynamic_assignment methods such as yolov7 or ' - 'rtmdet assigner, please load the checkpoint.') - - assert isinstance(model.bbox_head, (YOLOv5HeadAssigner, RTMHeadAssigner)),\ - 'Now, this script only support yolov5 and rtmdet, and ' \ + 'if you use dynamic_assignment methods such as YOLOv7 or ' + 'RTMDet assigner, please load the checkpoint.') + assert isinstance(model.bbox_head, (YOLOv5HeadAssigner, + YOLOv7HeadAssigner, + RTMHeadAssigner)), \ + 'Now, this script only support YOLOv5, YOLOv7 and RTMdet, and ' \ 'bbox_head must use ' \ - '`YOLOv5HeadAssigner or RTMHeadAssigner`. Please use `' \ + '`YOLOv5HeadAssigner or YOLOv7HeadAssigner or RTMHeadAssigner`.' \ + ' Please use `' \ 'yolov5_s-v61_syncbn_fast_8xb16-300e_coco_assignervisualization.py' \ + 'or yolov7_tiny_syncbn_fast_8x16b-300e_coco_assignervisualization.py' \ 'or rtmdet_s_syncbn_fast_8xb32-300e_coco_assignervisualization.py' \ """` as config file.""" model.eval() diff --git a/projects/assigner_visualization/configs/yolov7_tiny_syncbn_fast_8xb16-300e_coco_assignervisualization.py b/projects/assigner_visualization/configs/yolov7_tiny_syncbn_fast_8xb16-300e_coco_assignervisualization.py new file mode 100644 index 00000000..626dc18b --- /dev/null +++ b/projects/assigner_visualization/configs/yolov7_tiny_syncbn_fast_8xb16-300e_coco_assignervisualization.py @@ -0,0 +1,9 @@ +_base_ = ['../../../configs/yolov7/yolov7_tiny_syncbn_fast_8x16b-300e_coco.py'] + +custom_imports = dict(imports=[ + 'projects.assigner_visualization.detectors', + 'projects.assigner_visualization.dense_heads' +]) + +model = dict( + type='YOLODetectorAssigner', bbox_head=dict(type='YOLOv7HeadAssigner')) diff --git a/projects/assigner_visualization/dense_heads/__init__.py b/projects/assigner_visualization/dense_heads/__init__.py index fe41e5d6..e985d20c 100644 --- a/projects/assigner_visualization/dense_heads/__init__.py +++ b/projects/assigner_visualization/dense_heads/__init__.py @@ -1,5 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. from .rtmdet_head_assigner import RTMHeadAssigner from .yolov5_head_assigner import YOLOv5HeadAssigner +from .yolov7_head_assigner import YOLOv7HeadAssigner -__all__ = ['YOLOv5HeadAssigner', 'RTMHeadAssigner'] +__all__ = ['YOLOv5HeadAssigner', 'YOLOv7HeadAssigner', 'RTMHeadAssigner'] diff --git a/projects/assigner_visualization/dense_heads/yolov7_head_assigner.py b/projects/assigner_visualization/dense_heads/yolov7_head_assigner.py new file mode 100644 index 00000000..de2a90e3 --- /dev/null +++ b/projects/assigner_visualization/dense_heads/yolov7_head_assigner.py @@ -0,0 +1,159 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from typing import List, Union + +import torch +from mmdet.utils import InstanceList +from torch import Tensor + +from mmyolo.models import YOLOv7Head +from mmyolo.registry import MODELS + + +@MODELS.register_module() +class YOLOv7HeadAssigner(YOLOv7Head): + + def assign_by_gt_and_feat( + self, + cls_scores: List[Tensor], + bbox_preds: List[Tensor], + objectnesses: List[Tensor], + batch_gt_instances: InstanceList, + batch_img_metas: List[dict], + inputs_hw: Union[Tensor, tuple], + ) -> dict: + """Calculate the assigning results based on the gt and features + extracted by the detection head. + Args: + cls_scores (Sequence[Tensor]): Box scores for each scale level, + each is a 4D-tensor, the channel number is + num_priors * num_classes. + bbox_preds (Sequence[Tensor]): Box energies / deltas for each scale + level, each is a 4D-tensor, the channel number is + num_priors * 4. + objectnesses (Sequence[Tensor]): Score factor for + all scale level, each is a 4D-tensor, has shape + (batch_size, 1, H, W) + batch_gt_instances (list[:obj:`InstanceData`]): Batch of + gt_instance. It usually includes ``bboxes`` and ``labels`` + attributes. + batch_img_metas (list[dict]): Meta information of each image, e.g., + image size, scaling factor, etc. + inputs_hw (Union[Tensor, tuple]): Height and width of inputs size. + Returns: + dict[str, Tensor]: A dictionary of assigning results. + """ + device = cls_scores[0][0].device + + head_preds = self._merge_predict_results(bbox_preds, objectnesses, + cls_scores) + + batch_targets_normed = self._convert_gt_to_norm_format( + batch_gt_instances, batch_img_metas) + + # yolov5_assign and simota_assign + assigner_results = self.assigner( + head_preds, + batch_targets_normed, + batch_img_metas[0]['batch_input_shape'], + self.priors_base_sizes, + self.grid_offset, + near_neighbor_thr=self.near_neighbor_thr) + + # multi-level positive sample position. + mlvl_positive_infos = assigner_results['mlvl_positive_infos'] + # assigned results with label and bboxes information. + mlvl_targets_normed = assigner_results['mlvl_targets_normed'] + + assign_results = [] + for i in range(self.num_levels): + assign_results_feat = [] + # no gt bbox matches anchor + if mlvl_positive_infos[i].shape[0] == 0: + for k in range(self.num_base_priors): + assign_results_feat.append({ + 'stride': + self.featmap_strides[i], + 'grid_x_inds': + torch.zeros([0], dtype=torch.int64).to(device), + 'grid_y_inds': + torch.zeros([0], dtype=torch.int64).to(device), + 'img_inds': + torch.zeros([0], dtype=torch.int64).to(device), + 'class_inds': + torch.zeros([0], dtype=torch.int64).to(device), + 'retained_gt_inds': + torch.zeros([0], dtype=torch.int64).to(device), + 'prior_ind': + k + }) + assign_results.append(assign_results_feat) + continue + + # (batch_idx, prior_idx, x_scaled, y_scaled) + positive_info = mlvl_positive_infos[i] + targets_normed = mlvl_targets_normed[i] + priors_inds = positive_info[:, 1] + grid_x_inds = positive_info[:, 2] + grid_y_inds = positive_info[:, 3] + img_inds = targets_normed[:, 0] + class_inds = targets_normed[:, 1].long() + retained_gt_inds = self.get_gt_inds( + targets_normed, batch_targets_normed[0]).long() + for k in range(self.num_base_priors): + retained_inds = priors_inds == k + assign_results_prior = { + 'stride': self.featmap_strides[i], + 'grid_x_inds': grid_x_inds[retained_inds], + 'grid_y_inds': grid_y_inds[retained_inds], + 'img_inds': img_inds[retained_inds], + 'class_inds': class_inds[retained_inds], + 'retained_gt_inds': retained_gt_inds[retained_inds], + 'prior_ind': k + } + assign_results_feat.append(assign_results_prior) + assign_results.append(assign_results_feat) + return assign_results + + def get_gt_inds(self, assigned_target, gt_instance): + """Judging which one gt_ind is assigned by comparing assign_target and + origin target. + + Args: + assigned_target (Tensor(assign_nums,7)): YOLOv7 assigning results. + gt_instance (Tensor(gt_nums,7)): Normalized gt_instance, It + usually includes ``bboxes`` and ``labels`` attributes. + Returns: + gt_inds (Tensor): the index which one gt is assigned. + """ + gt_inds = torch.zeros(assigned_target.shape[0]) + for i in range(assigned_target.shape[0]): + gt_inds[i] = ((assigned_target[i] == gt_instance).sum( + dim=1) == 7).nonzero().squeeze() + return gt_inds + + def assign(self, batch_data_samples: Union[list, dict], + inputs_hw: Union[tuple, torch.Size]) -> dict: + """Calculate assigning results. + + This function is provided to the + `assigner_visualization.py` script. + Args: + batch_data_samples (List[:obj:`DetDataSample`], dict): The Data + Samples. It usually includes information such as + `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`. + inputs_hw: Height and width of inputs size + Returns: + dict: A dictionary of assigning components. + """ + if isinstance(batch_data_samples, list): + raise NotImplementedError( + 'assigning results_list is not implemented') + else: + # Fast version + cls_scores, bbox_preds, objectnesses = self( + batch_data_samples['feats']) + assign_inputs = (cls_scores, bbox_preds, objectnesses, + batch_data_samples['bboxes_labels'], + batch_data_samples['img_metas'], inputs_hw) + assign_results = self.assign_by_gt_and_feat(*assign_inputs) + return assign_results diff --git a/projects/assigner_visualization/detectors/yolo_detector_assigner.py b/projects/assigner_visualization/detectors/yolo_detector_assigner.py index edf0b828..65e6a1cf 100644 --- a/projects/assigner_visualization/detectors/yolo_detector_assigner.py +++ b/projects/assigner_visualization/detectors/yolo_detector_assigner.py @@ -3,7 +3,8 @@ from typing import Union from mmyolo.models import YOLODetector from mmyolo.registry import MODELS -from projects.assigner_visualization.dense_heads import RTMHeadAssigner +from projects.assigner_visualization.dense_heads import (RTMHeadAssigner, + YOLOv7HeadAssigner) @MODELS.register_module() @@ -23,7 +24,7 @@ class YOLODetectorAssigner(YOLODetector): assert isinstance(data, dict) assert len(data['inputs']) == 1, 'Only support batchsize == 1' data = self.data_preprocessor(data, True) - if isinstance(self.bbox_head, RTMHeadAssigner): + if isinstance(self.bbox_head, (YOLOv7HeadAssigner, RTMHeadAssigner)): data['data_samples']['feats'] = self.extract_feat(data['inputs']) inputs_hw = data['inputs'].shape[-2:] assign_results = self.bbox_head.assign(data['data_samples'], inputs_hw)