mirror of https://github.com/open-mmlab/mmyolo.git
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>pull/558/head
Zhaoyan Fang
GitHub
parent
6f0a765ea8
commit
d3179daa4d
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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'))
|
||||
|
|
@ -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']
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue