diff --git a/projects/assigner_visualization/dense_heads/rtmdet_head_assigner.py b/projects/assigner_visualization/dense_heads/rtmdet_head_assigner.py index ae0e4651..b5003eef 100644 --- a/projects/assigner_visualization/dense_heads/rtmdet_head_assigner.py +++ b/projects/assigner_visualization/dense_heads/rtmdet_head_assigner.py @@ -7,6 +7,7 @@ from mmdet.utils import InstanceList from torch import Tensor from mmyolo.models import RTMDetHead +from mmyolo.models.utils import gt_instances_preprocess from mmyolo.registry import MODELS @@ -42,8 +43,9 @@ class RTMHeadAssigner(RTMDetHead): num_imgs = len(batch_img_metas) featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores] assert len(featmap_sizes) == self.prior_generator.num_levels + prior_offset = self.prior_generator.offset - gt_info = self.gt_instances_preprocess(batch_gt_instances, num_imgs) + gt_info = gt_instances_preprocess(batch_gt_instances, num_imgs) gt_labels = gt_info[:, :, :1] gt_bboxes = gt_info[:, :, 1:] # xyxy pad_bbox_flag = (gt_bboxes.sum(-1, keepdim=True) > 0).float() @@ -122,7 +124,9 @@ class RTMHeadAssigner(RTMDetHead): 'retained_gt_inds': torch.zeros([0], dtype=torch.int64).to(device), 'prior_ind': - 0 + 0, + 'offset': + prior_offset } else: w = inputs_hw[1] // self.featmap_strides[i] @@ -137,7 +141,8 @@ class RTMHeadAssigner(RTMDetHead): 'img_inds': img_inds[retained_inds], 'class_inds': labels[retained_inds], 'retained_gt_inds': matched_gt_inds[retained_inds], - 'prior_ind': 0 + 'prior_ind': 0, + 'offset': prior_offset } assign_results.append([assign_results_prior]) return assign_results diff --git a/projects/assigner_visualization/visualization/assigner_visualizer.py b/projects/assigner_visualization/visualization/assigner_visualizer.py index 73aba9f8..41d8f62d 100644 --- a/projects/assigner_visualization/visualization/assigner_visualizer.py +++ b/projects/assigner_visualization/visualization/assigner_visualizer.py @@ -274,9 +274,11 @@ class YOLOAssignerVisualizer(DetLocalVisualizer): grid_y_inds = assign_results_prior['grid_y_inds'] class_inds = assign_results_prior['class_inds'] prior_ind = assign_results_prior['prior_ind'] + offset = assign_results_prior.get('offset', 0.5) + if show_prior: self.draw_prior(grid_x_inds, grid_y_inds, class_inds, - stride, feat_ind, prior_ind) + stride, feat_ind, prior_ind, offset) # draw matched gt retained_gt_inds = assign_results_prior['retained_gt_inds'] @@ -286,7 +288,7 @@ class YOLOAssignerVisualizer(DetLocalVisualizer): # draw positive self.draw_positive_assign(grid_x_inds, grid_y_inds, class_inds, stride, gt_instances.bboxes, - retained_gt_inds) + retained_gt_inds, offset) # draw title if self.priors_size is not None: @@ -319,4 +321,7 @@ class YOLOAssignerVisualizer(DetLocalVisualizer): img_show_list.append(np.concatenate(img_show_list_feat, axis=1)) # Merge all images into one image - return np.concatenate(img_show_list, axis=0) + h, w = img_show.shape[:2] + num_priors_per_feat = img_show_list[0].shape[1] // w + axis = 0 if num_priors_per_feat > 1 else 1 + return np.concatenate(img_show_list, axis=axis)