From db15804710ee0585b04384254d5cee97ae1028f0 Mon Sep 17 00:00:00 2001
From: Nathan Michlo <NathanJMichlo@gmail.com>
Date: Tue, 9 Nov 2021 11:39:40 +0200
Subject: [PATCH] cleanup new function

---
 dataset_tool.py | 44 +++++++++++++++++++-------------------------
 1 file changed, 19 insertions(+), 25 deletions(-)

diff --git a/dataset_tool.py b/dataset_tool.py
index b241faa0..94dbe89f 100644
--- a/dataset_tool.py
+++ b/dataset_tool.py
@@ -250,8 +250,8 @@ def center_crop_wide(width, height, img):
         canvas[(width - height) // 2 : (width + height) // 2, :] = img
         return canvas
 
-    def resize_pad(width, height, img):
-        # fix dims
+    def _to_rgb(img):
+        # add channel dim
         if img.ndim == 2:
             img = img[:, :, None]
         assert img.ndim == 3, f'input image has incorrect number of dimensions, required 2 (H, W) or 3 (H, W, C), got: {img.shape}'
@@ -261,34 +261,28 @@ def resize_pad(width, height, img):
         elif img.shape[-1] == 4:
             img = img[:, :, :3]
         assert img.shape[-1] == 3, f'input image must have 1 or 3 channels, got: {img.shape}'
+        return img
+
+    def resize_pad(width, height, img):
+        img = _to_rgb(img)
         # exit early
-        h, w, c = img.shape
-        if width == w and height == h:
+        img_h, img_w = img.shape[:2]
+        if width == img_w and height == img_h:
             return img
-        # get scale size
-        rh = height / h
-        rw = width / w
-        sh = int(round(h * min(rh, rw), 5))  # avoid precision errors
-        sw = int(round(w * min(rh, rw), 5))  # avoid precision errors
-        assert sh <= height and sw <= height, f'scaled shape {nw}x{nh} is not smaller than or equal to the required shape: {width}x{height} this is a bug:'
+        # get scale size, avoiding precision errors
+        scale_ratio = max(img_h / height, img_w / width)
+        scale_h = int(round(img_h / scale_ratio, 5))
+        scale_w = int(round(img_w / scale_ratio, 5))
+        assert scale_h <= height and scale_w <= width, f'scaled shape {scale_w}x{scale_h} is not smaller than or equal to the required shape: {width}x{height} this is a bug:'
         # scale image
-        img = scale(sw, sh, img)
-        nh, nw, nc = img.shape
-        assert nh == sh and nw == sw, f'scaled shape {nw}x{nh} does not match required scaled shape: {sw}x{sh} this is a bug!'
+        img = scale(scale_w, scale_h, img)
         # pad the image if needed
-        ph = height - sh
-        pw = width - sw
-        if ph != 0 or pw != 0:
-            assert ph >= 0 and pw >= 0, f'target width={repr(width)} height={repr(width)}, pad amount: pw={repr(pw)} ph={repr(ph)}'
-            img = np.pad(img, (
-                (int(np.floor(ph/2)), int(np.ceil(ph/2))),
-                (int(np.floor(pw/2)), int(np.ceil(pw/2))),
-                (0, 0),
-            ))
+        pad_h, pad_w = height - scale_h, width - scale_w
+        if pad_h != 0 or pad_w != 0:
+            pad_dims = [[np.floor(pad_h/2), np.ceil(pad_h/2)], [np.floor(pad_w/2), np.ceil(pad_w/2)], [0, 0]]  # (H,W,C)
+            img = np.pad(img, np.array(pad_dims).astype('int'))
         # check the shape
-        oh, ow, oc = img.shape
-        # print((h, w, c), h/w, '->', (nh, nw, nc), nh/nw, '->', (oh, ow, oc), oh/ow)
-        assert oh == height and ow == width, f'output shape {ow}x{oh} does not match required shape: {width}x{height} this is a bug!'
+        assert img.shape[0] == height and img.shape[1] == width, f'output shape {img.shape[1]}x{img.shape[0]} does not match required shape: {width}x{height} this is a bug!'
         return img
 
     if transform is None: