-
Notifications
You must be signed in to change notification settings - Fork 7
/
Copy pathinference.py
49 lines (41 loc) · 2 KB
/
inference.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import torch
from PIL import Image
from utils import *
import torch.nn.functional as F
import numpy as np
def get_3angle(image, dino, val_preprocess, device):
# 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+180], dim=-1)
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
angles[2] = gaus_ro_pred - 90
angles[3] = confidence
return angles
def get_3angle_infer_aug(origin_img, rm_bkg_img, dino, val_preprocess, device):
# image = Image.open(image_path).convert('RGB')
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():
dino_pred = dino(image_inputs)
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+180], 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
angles[2] = gaus_ro_pred - 90
angles[3] = confidence
return angles