fix pad2D potential error (#250)

tusen-ai · Oct 24, 2019 · f09f507 · f09f507
1 parent 67c79f5
commit f09f507
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/core/detection_input.py b/core/detection_input.py
@@ -304,8 +304,9 @@ def apply(self, input_record):
         image = input_record["image"]
         gt_bbox = input_record["gt_bbox"]
 
+        origin_h, origin_w = input_record["h"], input_record["w"]
         h, w = image.shape[:2]
-        shape = (p.long, p.short, 3) if h >= w \
+        shape = (p.long, p.short, 3) if origin_h >= origin_w \
             else (p.short, p.long, 3)
 
         padded_image = np.zeros(shape, dtype=np.float32)
@@ -334,8 +335,9 @@ def apply(self, input_record):
 
         image = input_record["image"]
 
+        origin_h, origin_w = input_record["h"], input_record["w"]
         h, w = image.shape[:2]
-        shape = (p.long, p.short, 3) if h >= w \
+        shape = (p.long, p.short, 3) if origin_h >= origin_w \
             else (p.short, p.long, 3)
 
         padded_image = np.zeros(shape, dtype=np.float32)