infer aug

app.py

@@ -8,6 +8,7 @@ import os
 import matplotlib.pyplot as plt
 import io
 from PIL import Image
+import random
 import rembg
 from typing import Any
 import torch.nn.functional as F
@@ -97,6 +98,37 @@ def remove_background(image: Image,
         image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
     return image
 
+def random_crop(image, crop_scale=(0.8, 0.95)):
+    """
+    随机裁切图片
+        image (numpy.ndarray):  (H, W, C)。
+        crop_scale (tuple): (min_scale, max_scale)。
+    """
+    assert isinstance(image, Image.Image), "iput must be PIL.Image.Image"
+    assert len(crop_scale) == 2 and 0 < crop_scale[0] <= crop_scale[1] <= 1
+
+    width, height = image.size
+
+    # 计算裁切的高度和宽度
+    crop_width = random.randint(int(width * crop_scale[0]), int(width * crop_scale[1]))
+    crop_height = random.randint(int(height * crop_scale[0]), int(height * crop_scale[1]))
+
+    # 随机选择裁切的起始点
+    left = random.randint(0, width - crop_width)
+    top = random.randint(0, height - crop_height)
+
+    # 裁切图片
+    cropped_image = image.crop((left, top, left + crop_width, top + crop_height))
+
+    return cropped_image
+
+def get_crop_images(img, num=3):
+    cropped_images = []
+    for i in range(num):
+        cropped_images.append(random_crop(img))
+    return cropped_images
+
+
 def get_3angle(image):
     
     # image = Image.open(image_path).convert('RGB')
@@ -108,7 +140,7 @@ def get_3angle(image):
     gaus_ax_pred   = torch.argmax(dino_pred[:, 0:360], dim=-1)
     gaus_pl_pred   = torch.argmax(dino_pred[:, 360:360+180], dim=-1)
     gaus_ro_pred   = torch.argmax(dino_pred[:, 360+180:360+180+60], dim=-1)
-    confidence     = F.softmax(dino_pred[:, -2:], dim=-1)[0]
+    confidence     = F.softmax(dino_pred[:, -2:], dim=-1)[0][0]
     angles = torch.zeros(4)
     angles[0]  = gaus_ax_pred
     angles[1]  = gaus_pl_pred - 90
@@ -116,18 +148,86 @@ def get_3angle(image):
     angles[3]  = confidence
     return angles
 
+def remove_outliers_and_average(tensor, threshold=1.5):
+    assert tensor.dim() == 1, "dimension of input Tensor must equal to 1"
+
+    q1 = torch.quantile(tensor, 0.25)
+    q3 = torch.quantile(tensor, 0.75)
+    iqr = q3 - q1
+
+    lower_bound = q1 - threshold * iqr
+    upper_bound = q3 + threshold * iqr
+
+    non_outliers = tensor[(tensor >= lower_bound) & (tensor <= upper_bound)]
+
+    if len(non_outliers) == 0:
+        return tensor.mean().item()
+
+    return non_outliers.mean().item()
+
+
+def remove_outliers_and_average_circular(tensor, threshold=1.5):
+    assert tensor.dim() == 1, "dimension of input Tensor must equal to 1"
+
+    # 将角度转换为二维平面上的点
+    radians = tensor * torch.pi / 180.0
+    x_coords = torch.cos(radians)
+    y_coords = torch.sin(radians)
+
+    # 计算平均向量
+    mean_x = torch.mean(x_coords)
+    mean_y = torch.mean(y_coords)
+
+    differences = torch.sqrt((x_coords - mean_x) * (x_coords - mean_x) + (y_coords - mean_y) * (y_coords - mean_y))
+
+    # 计算四分位数和 IQR
+    q1 = torch.quantile(differences, 0.25)
+    q3 = torch.quantile(differences, 0.75)
+    iqr = q3 - q1
+
+    # 计算上下限
+    lower_bound = q1 - threshold * iqr
+    upper_bound = q3 + threshold * iqr
+
+    # 筛选非离群点
+    non_outliers = tensor[(differences >= lower_bound) & (differences <= upper_bound)]
+
+    if len(non_outliers) == 0:
+        mean_angle = torch.atan2(mean_y, mean_x) * 180.0 / torch.pi
+        mean_angle = (mean_angle + 360) % 360
+        return mean_angle  # 如果没有非离群点，返回 None
+
+    # 对非离群点再次计算平均向量
+    radians = non_outliers * torch.pi / 180.0
+    x_coords = torch.cos(radians)
+    y_coords = torch.sin(radians)
+
+    mean_x = torch.mean(x_coords)
+    mean_y = torch.mean(y_coords)
+
+    mean_angle = torch.atan2(mean_y, mean_x) * 180.0 / torch.pi
+    mean_angle = (mean_angle + 360) % 360
+
+    return mean_angle
+
 def get_3angle_infer_aug(image):
     
     # image = Image.open(image_path).convert('RGB')
+    image = get_crop_images(image, num=6)
     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():
         dino_pred = dino(image_inputs)
 
-    gaus_ax_pred   = torch.argmax(dino_pred[:, 0:360], dim=-1)
-    gaus_pl_pred   = torch.argmax(dino_pred[:, 360:360+180], dim=-1)
-    gaus_ro_pred   = torch.argmax(dino_pred[:, 360+180:360+180+60], dim=-1)
-    confidence     = F.softmax(dino_pred[:, -2:], dim=-1)[0]
+    gaus_ax_pred   = torch.argmax(dino_pred[:, 0:360], dim=-1).to(torch.float32)
+    gaus_pl_pred   = torch.argmax(dino_pred[:, 360:360+180], dim=-1).to(torch.float32)
+    gaus_ro_pred   = torch.argmax(dino_pred[:, 360+180:360+180+60], dim=-1).to(torch.float32)
+    
+    gaus_ax_pred   = remove_outliers_and_average_circular(gaus_ax_pred)
+    gaus_pl_pred   = remove_outliers_and_average(gaus_pl_pred)
+    gaus_ro_pred   = remove_outliers_and_average(gaus_ro_pred)
+    
+    confidence     = torch.mean(F.softmax(dino_pred[:, -2:], dim=-1), dim=0)[0]
     angles = torch.zeros(4)
     angles[0]  = gaus_ax_pred
     angles[1]  = gaus_pl_pred - 90
@@ -221,7 +321,7 @@ def infer_func(img, do_rm_bkg, do_infer_aug):
     
     res_img = figure_to_img(fig)
     # axis_model = "axis.obj"
-    return [res_img, float(angles[0]), float(angles[1]), float(angles[2]), float(angles[3])]
+    return [res_img, round(float(angles[0]), 2), round(float(angles[1]), 2), round(float(angles[2]), 2), round(float(angles[3]), 2)]
 
 server = gr.Interface(
     flagging_mode='never',