infer aug + remove bkg

app.py

@@ -210,10 +210,10 @@ def remove_outliers_and_average_circular(tensor, threshold=1.5):
 
     return mean_angle
 
-def get_3angle_infer_aug(image):
+def get_3angle_infer_aug(origin_img, rm_bkg_img):
     
     # image = Image.open(image_path).convert('RGB')
-    image = get_crop_images(image, num=6)
+    image = get_crop_images(origin_img, num=3) + get_crop_images(rm_bkg_img, num=3)
     image_inputs = val_preprocess(images = image)
     image_inputs['pixel_values'] = torch.from_numpy(np.array(image_inputs['pixel_values'])).to(device)
     with torch.no_grad():
@@ -267,21 +267,22 @@ def figure_to_img(fig):
     return image
 
 def infer_func(img, do_rm_bkg, do_infer_aug):
-    img = Image.fromarray(img)
-    img = background_preprocess(img, do_rm_bkg)
+    origin_img = Image.fromarray(img)
     if do_infer_aug:
-        angles = get_3angle_infer_aug(img)
+        rm_bkg_img = background_preprocess(origin_img, True)
+        angles = get_3angle_infer_aug(origin_img, rm_bkg_img)
     else:
-        angles = get_3angle(img)
+        rm_bkg_img = background_preprocess(origin_img, do_rm_bkg)
+        angles = get_3angle(rm_bkg_img)
     
     fig, ax = plt.subplots(figsize=(8, 8))
 
-    w, h = img.size
+    w, h = rm_bkg_img.size
     if h>w:
         extent = [-5*w/h, 5*w/h, -5, 5]
     else:
         extent = [-5, 5, -5*h/w, 5*h/w]
-    ax.imshow(img, extent=extent, zorder=0, aspect ='auto')  # extent 设置图片的显示范围
+    ax.imshow(rm_bkg_img, extent=extent, zorder=0, aspect ='auto')  # extent 设置图片的显示范围
 
     origin = np.array([0, 0])