confidence added

app.py

@@ -10,6 +10,7 @@ import io
 from PIL import Image
 import rembg
 from typing import Any
+import torch.nn.functional as F
 
 
 from huggingface_hub import hf_hub_download
@@ -107,11 +108,31 @@ 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)
-    angles = torch.zeros(3)
+    confidence     = F.softmax(dino_pred[:, -2:], dim=-1)[0]
+    angles = torch.zeros(4)
     angles[0]  = gaus_ax_pred
     angles[1]  = gaus_pl_pred - 90
     angles[2]  = gaus_ro_pred - 30
+    angles[3]  = confidence
+    return angles
+
+def get_3angle_infer_aug(image):
     
+    # image = Image.open(image_path).convert('RGB')
+    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]
+    angles = torch.zeros(4)
+    angles[0]  = gaus_ax_pred
+    angles[1]  = gaus_pl_pred - 90
+    angles[2]  = gaus_ro_pred - 30
+    angles[3]  = confidence
     return angles
 
 def scale(x):
@@ -145,10 +166,13 @@ def figure_to_img(fig):
         image = Image.open(buf).copy()
     return image
 
-def infer_func(img, do_rm_bkg):
+def infer_func(img, do_rm_bkg, do_infer_aug):
     img = Image.fromarray(img)
     img = background_preprocess(img, do_rm_bkg)
-    angles = get_3angle(img)
+    if do_infer_aug:
+        angles = get_3angle_infer_aug(img)
+    else:
+        angles = get_3angle(img)
     
     fig, ax = plt.subplots(figsize=(8, 8))
 
@@ -197,21 +221,23 @@ def infer_func(img, do_rm_bkg):
     
     res_img = figure_to_img(fig)
     # axis_model = "axis.obj"
-    return [res_img, float(angles[0]), float(angles[1]), float(angles[2])]
+    return [res_img, float(angles[0]), float(angles[1]), float(angles[2]), float(angles[3])]
 
 server = gr.Interface(
     flagging_mode='never',
     fn=infer_func, 
     inputs=[
         gr.Image(height=512, width=512, label="upload your image"),
-        gr.Checkbox(label="Remove Background", value=True)
+        gr.Checkbox(label="Remove Background", value=True),
+        gr.Checkbox(label="Inference time augmentation", value=False)
     ], 
     outputs=[
         gr.Image(height=512, width=512, label="result image"),
         # gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model"),
         gr.Textbox(lines=1, label='Azimuth(0~360°)'),
         gr.Textbox(lines=1, label='Polar(-90~90°)'),
-        gr.Textbox(lines=1, label='Rotation(-90~90°)')
+        gr.Textbox(lines=1, label='Rotation(-90~90°)'),
+        gr.Textbox(lines=1, label='Confidence(0~1)')
     ]
 )