From a0ba64ee18f2d8cde9f7e00938aec0c988856990 Mon Sep 17 00:00:00 2001
From: Bin Long <sjtubinlong@users.noreply.github.com>
Date: Wed, 18 Sep 2019 18:05:14 +0800
Subject: [PATCH] Update visualize.py

---
 inference/tools/visualize.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/inference/tools/visualize.py b/inference/tools/visualize.py
index 87373e91..e0740d46 100644
--- a/inference/tools/visualize.py
+++ b/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)
-- 
GitLab