Update visualize.py

inference/tools/visualize.py

@@ -10,10 +10,9 @@ color_map = [[128, 64, 128], [244, 35, 231], [69, 69, 69], [102, 102, 156],
 
 im = cv2.imread(sys.argv[1])
 
-# Please note (224, 224) just for daheng model
 print("visualizing...")
-for i in range(0, 224):
-    for j in range(0, 224):
+for i in range(0, im.rows):
+    for j in range(0, im.cols):
         im[i, j] = color_map[im[i, j, 0]]
 
 cv2.imwrite(sys.argv[1], im)