Fix visualization

Author: ruotianluo

lib/nets/network.py

@@ -8,6 +8,7 @@
 from __future__ import print_function
 
 import math
+import numpy as np
 
 import torch
 import torch.nn as nn
@@ -49,9 +50,9 @@ def __init__(self):
   def _add_gt_image(self):
     # add back mean
     image = self._image_gt_summaries['image'] + cfg.PIXEL_MEANS
-    resized = imresize(image, self._im_info[:2] / self._im_info[2])
+    image = imresize(image[0], self._im_info[:2] / self._im_info[2])
     # BGR to RGB (opencv uses BGR)
-    self._gt_image = image[:,:,::-1].copy(order='C')
+    self._gt_image = image[np.newaxis, :,:,::-1].copy(order='C')
 
   def _add_gt_image_summary(self):
     # use a customized visualization function to visualize the boxes

lib/utils/visualization.py

@@ -75,16 +75,15 @@ def draw_bounding_boxes(image, gt_boxes, im_info):
   disp_image = Image.fromarray(np.uint8(image[0]))
 
   for i in range(num_boxes):
-    this_class = int(gt_boxes[i, 4])
-    disp_image = draw_single_box(disp_image, 
-                                gt_boxes[i, 0],
-                                gt_boxes[i, 1],
-                                gt_boxes[i, 2],
-                                gt_boxes[i, 3],
-                                'N%d-C%d' % (i, this_class),
-                                font,
-                                color=STANDARD_COLORS[this_class % num_colors])
+    this_class = int(gt_boxes_new[i, 4])
+    disp_image = _draw_single_box(disp_image, 
+                                gt_boxes_new[i, 0],
+                                gt_boxes_new[i, 1],
+                                gt_boxes_new[i, 2],
+                                gt_boxes_new[i, 3],
+                                'N%02d-C%02d' % (i, this_class),
+                                FONT,
+                                color=STANDARD_COLORS[this_class % NUM_COLORS])
 
   image[0, :] = np.array(disp_image)
-  return image
-
+  return image