Add rescale function

app.py

@@ -21,7 +21,7 @@ def filesort(img, model):
     img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     h, w = img.shape
     img_out = preprocessing(img, model)
-    return img_out, h, w, img, ori
+    return img_out, h, w, ori
 
 def preprocessing(img, model='SE-RegUNet 4GF'):
     # print(img.shape, img.dtype)
@@ -46,7 +46,14 @@ def preprocessing(img, model='SE-RegUNet 4GF'):
         img_out = np.stack((image1,)*3, axis=0)
     return img_out
 
-def process_input_image(img, model, rescale):
+def inference(pipe, img, model):
+    with torch.no_grad():
+        if model == 'AngioNet':
+            img = torch.cat([img, img], dim=0)
+        logit = np.round(torch.softmax(pipe.forward(img), dim=1).detach().cpu().numpy()[0, 0]).astype(np.uint8)   
+    return logit
+
+def process_input_image(img, model, scale):
     ori_img = img.copy()
     h, w, _ = ori_img.shape
     pad_h = h % 32
@@ -72,13 +79,21 @@ def process_input_image(img, model, rescale):
         pipe = torch.jit.load('./model/UNet3plus.pt')
     pipe = pipe.to(device).eval()
     
+    scale = int(rescale.split('x')[0])
+    
     start = time.time()
-    img, h, w, ori_gray, ori = filesort(img, model)
-    img = torch.FloatTensor(img).unsqueeze(0).to(device)
-    with torch.no_grad():
-        if model == 'AngioNet':
-            img = torch.cat([img, img], dim=0)
-        logit = np.round(torch.softmax(pipe.forward(img), dim=1).detach().cpu().numpy()[0, 0]).astype(np.uint8)
+    if scale == 1:
+        img, h, w, ori = filesort(img, model)
+        img = torch.FloatTensor(img).unsqueeze(0).to(device)
+        logit = inference(pipe, img, model)
+    else:
+        len_h, len_w = img.shape[0] // scale, img.shape[1] // scale
+        logit = np.zeros(img.shape, np.float32)
+        for x in range(scale):
+            for y in range(scale):
+                temp_img, _, _, _ = filesort(img[len_h * x : len_h * (x + 1), len_w * y : len_w * (y + 1)])
+                temp_img = torch.FloatTensor(temp_img).unsqueeze(0).to(device)
+                logit[len_h * x : len_h * (x + 1), len_w * y : len_w * (y + 1)] = inference(pipe, temp_img, model)
     spent = time.time() - start
     spent = f"{spent:.3f} seconds"
     
@@ -108,7 +123,7 @@ with my_app:
                     img_source = gr.Image(label="Please select angiogram.", value='./example/angio.png', shape=(512, 512))
                     model_choice = gr.Dropdown(['SE-RegUNet 4GF', 'SE-RegUNet 16GF', 'AngioNet', 'EffUNet++ B5', 
                                                 'Reg-SA-UNet++', 'UNet3+'], label='Model', info='Which model to infer?')
-                    model_rescale = gr.Dropdown(['2x2', '4x4', '8x8', '16x16'], label='Rescale', info='How many batches?')
+                    model_rescale = gr.Dropdown(['1x1', '2x2', '4x4', '8x8', '16x16'], label='Rescale', info='How many batches?')
                     source_image_loader = gr.Button("Vessel Segment")
                 with gr.Column():
                     time_spent = gr.Label(label="Time Spent (Preprocessing + Inference)")