[Fix] fix visualization in sot demo (#213)

* [Fix] fix visualization in sot demo

* fix some bugs

* fix docstring

* fix a typo

demo/demo_sot.py

@@ -54,17 +54,18 @@ def main():
         result = inference_sot(model, frame, init_bbox, frame_id)
 
         track_bbox = result['bbox']
-        cv2.rectangle(
-            frame, (track_bbox[0], track_bbox[1]),
-            (track_bbox[2], track_bbox[3]),
-            args.color,
-            thickness=args.thickness)
+        vis_frame = model.show_result(
+            frame,
+            track_bbox,
+            color=args.color,
+            thickness=args.thickness,
+            show=False)
 
         if save_out_video:
-            videoWriter.write(frame)
+            videoWriter.write(vis_frame)
 
         if args.show:
-            cv2.imshow(args.input, frame)
+            cv2.imshow(args.input, vis_frame)
 
         if cv2.waitKey(1) & 0xFF == ord('q'):
             break

mmtrack/models/sot/base.py

@@ -1,6 +1,8 @@
 from abc import ABCMeta, abstractmethod
 from collections import OrderedDict
 
+import mmcv
+import numpy as np
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -232,3 +234,46 @@ def val_step(self, data, optimizer):
             loss=loss, log_vars=log_vars, num_samples=len(data['img_metas']))
 
         return outputs
+
+    def show_result(self,
+                    img,
+                    result,
+                    color='green',
+                    thickness=1,
+                    show=False,
+                    win_name='',
+                    wait_time=0,
+                    out_file=None):
+        """Visualize tracking results.
+
+        Args:
+            img (str or ndarray): The image to be displayed.
+            result (ndarray): ndarray of shape (4, ).
+            color (str or tuple or Color, optional): color of bbox.
+                Defaults to green.
+            thickness (int, optional): Thickness of lines.
+                Defaults to 1.
+            show (bool, optional): Whether to show the image.
+                Defaults to False.
+            win_name (str, optional): The window name.
+                Defaults to ''.
+            wait_time (int, optional): Value of waitKey param.
+                Defaults to 0.
+            out_file (str, optional): The filename to write the image.
+                Defaults to None.
+
+        Returns:
+            ndarray: Visualized image.
+        """
+        assert result.ndim == 1
+        assert result.shape[0] == 4
+        mmcv.imshow_bboxes(
+            img,
+            result[np.newaxis, :],
+            colors=color,
+            thickness=thickness,
+            show=show,
+            win_name=win_name,
+            wait_time=wait_time,
+            out_file=out_file)
+        return img