Fix RTMDet visualization to new code (#550)

* fix bug to newer code

* adjust picture format

* add offset for rtmdet

* rewrite axis to fix format

* easier way to get axis

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

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)