소스 검색

update PreProcessor

yjh0410 1 년 전
부모
커밋
8015d05b19
1개의 변경된 파일25개의 추가작업 그리고 15개의 파일을 삭제
  1. 25 15
      utils/misc.py

+ 25 - 15
utils/misc.py

@@ -330,8 +330,9 @@ def multiclass_nms(scores, labels, bboxes, nms_thresh, num_classes, class_agnost
 # ---------------------------- Processor for Deployment ----------------------------
 ## Pre-processer
 class PreProcessor(object):
-    def __init__(self, img_size):
+    def __init__(self, img_size, keep_ratio=True):
         self.img_size = img_size
+        self.keep_ratio = keep_ratio
         self.input_size = [img_size, img_size]
         
 
@@ -346,23 +347,32 @@ class PreProcessor(object):
         else:
             padded_img = np.ones(self.input_size, np.float32) * 114.
         # resize
-        orig_h, orig_w = image.shape[:2]
-        r = min(self.input_size[0] / orig_h, self.input_size[1] / orig_w)
-        resize_size = (int(orig_w * r), int(orig_h * r))
-        if r != 1:
-            resized_img = cv2.resize(image, resize_size, interpolation=cv2.INTER_LINEAR)
-        else:
-            resized_img = image
+        if self.keep_ratio:
+            orig_h, orig_w = image.shape[:2]
+            r = min(self.input_size[0] / orig_h, self.input_size[1] / orig_w)
+            resize_size = (int(orig_w * r), int(orig_h * r))
+            if r != 1:
+                resized_img = cv2.resize(image, resize_size, interpolation=cv2.INTER_LINEAR)
+            else:
+                resized_img = image
 
-        # padding
-        padded_img[:resized_img.shape[0], :resized_img.shape[1]] = resized_img
-        
-        # [H, W, C] -> [C, H, W]
-        padded_img = padded_img.transpose(swap)
-        padded_img = np.ascontiguousarray(padded_img, dtype=np.float32) / 255.
+            # padding
+            padded_img[:resized_img.shape[0], :resized_img.shape[1]] = resized_img
+            
+            # [H, W, C] -> [C, H, W]
+            padded_img = padded_img.transpose(swap)
+            padded_img = np.ascontiguousarray(padded_img, dtype=np.float32) / 255.
 
+            return padded_img, r
+        else:
+            orig_h, orig_w = image.shape[:2]
+            r = np.array([self.input_size[0] / orig_w, self.input_size[1] / orig_w])
+            if [orig_h, orig_w] == self.input_size:
+                resized_img = image
+            else:
+                resized_img = cv2.resize(image, self.input_size, interpolation=cv2.INTER_LINEAR)
 
-        return padded_img, r
+            return resized_img, r
 
 ## Post-processer
 class PostProcessor(object):