Initial Submit

app.py

@@ -51,7 +51,12 @@ def process_input_image(img, model, rescale):
     h, w, _ = ori_img.shape
     pad_h = h % 32
     pad_w = w % 32
-    img = ori_img[pad_h//2:-pad_h//2, pad_w//2:-pad_w//2]
+    if pad_h == 0 and pad_w > 0:
+        img = ori_img[:, pad_w//2:-pad_w//2]
+    elif pad_h > 0 and pad_w == 0:
+        img = ori_img[pad_h//2:-pad_h//2, :]
+    if pad_h == 0 and pad_w == 0:
+        img = ori_img[pad_h//2:-pad_h//2, pad_w//2:-pad_w//2]
     
     if model == 'SE-RegUNet 4GF':
         pipe = torch.jit.load('./model/SERegUNet4GF.pt')
@@ -81,7 +86,12 @@ def process_input_image(img, model, rescale):
     # img_out = cv2.cvtColor(ori, cv2.COLOR_GRAY2RGB)
     img_out = ori.copy()
     img_out[logit, 0] = 255
-    ori_img[pad_h//2:-pad_h//2, pad_w//2:-pad_w//2] = img_out
+    if pad_h == 0 and pad_w > 0:
+        ori_img[:, pad_w//2:-pad_w//2] = img_out
+    elif pad_h > 0 and pad_w == 0:
+        ori_img[pad_h//2:-pad_h//2, :] = img_out
+    if pad_h == 0 and pad_w == 0:
+        ori_img[pad_h//2:-pad_h//2, pad_w//2:-pad_w//2] = img_out
     return spent, ori_img