Spaces:
Running
Running
zhang-ziang
commited on
Commit
·
00e3bbc
1
Parent(s):
60644d7
update model weight
Browse files- app.py +2 -2
- inference.py +4 -4
app.py
CHANGED
@@ -8,7 +8,7 @@ from inference import *
|
|
8 |
from utils import *
|
9 |
|
10 |
from huggingface_hub import hf_hub_download
|
11 |
-
ckpt_path = hf_hub_download(repo_id="Viglong/Orient-Anything", filename="
|
12 |
print(ckpt_path)
|
13 |
|
14 |
save_path = './'
|
@@ -16,7 +16,7 @@ device = 'cpu'
|
|
16 |
dino = DINOv2_MLP(
|
17 |
dino_mode = 'large',
|
18 |
in_dim = 1024,
|
19 |
-
out_dim = 360+180+
|
20 |
evaluate = True,
|
21 |
mask_dino = False,
|
22 |
frozen_back = False
|
|
|
8 |
from utils import *
|
9 |
|
10 |
from huggingface_hub import hf_hub_download
|
11 |
+
ckpt_path = hf_hub_download(repo_id="Viglong/Orient-Anything", filename="ronormsigma1/dino_weight.pt", repo_type="model", cache_dir='./', resume_download=True)
|
12 |
print(ckpt_path)
|
13 |
|
14 |
save_path = './'
|
|
|
16 |
dino = DINOv2_MLP(
|
17 |
dino_mode = 'large',
|
18 |
in_dim = 1024,
|
19 |
+
out_dim = 360+180+360+2,
|
20 |
evaluate = True,
|
21 |
mask_dino = False,
|
22 |
frozen_back = False
|
inference.py
CHANGED
@@ -14,12 +14,12 @@ def get_3angle(image, dino, val_preprocess, device):
|
|
14 |
|
15 |
gaus_ax_pred = torch.argmax(dino_pred[:, 0:360], dim=-1)
|
16 |
gaus_pl_pred = torch.argmax(dino_pred[:, 360:360+180], dim=-1)
|
17 |
-
gaus_ro_pred = torch.argmax(dino_pred[:, 360+180:360+180+
|
18 |
confidence = F.softmax(dino_pred[:, -2:], dim=-1)[0][0]
|
19 |
angles = torch.zeros(4)
|
20 |
angles[0] = gaus_ax_pred
|
21 |
angles[1] = gaus_pl_pred - 90
|
22 |
-
angles[2] = gaus_ro_pred -
|
23 |
angles[3] = confidence
|
24 |
return angles
|
25 |
|
@@ -34,7 +34,7 @@ def get_3angle_infer_aug(origin_img, rm_bkg_img, dino, val_preprocess, device):
|
|
34 |
|
35 |
gaus_ax_pred = torch.argmax(dino_pred[:, 0:360], dim=-1).to(torch.float32)
|
36 |
gaus_pl_pred = torch.argmax(dino_pred[:, 360:360+180], dim=-1).to(torch.float32)
|
37 |
-
gaus_ro_pred = torch.argmax(dino_pred[:, 360+180:360+180+
|
38 |
|
39 |
gaus_ax_pred = remove_outliers_and_average_circular(gaus_ax_pred)
|
40 |
gaus_pl_pred = remove_outliers_and_average(gaus_pl_pred)
|
@@ -44,6 +44,6 @@ def get_3angle_infer_aug(origin_img, rm_bkg_img, dino, val_preprocess, device):
|
|
44 |
angles = torch.zeros(4)
|
45 |
angles[0] = gaus_ax_pred
|
46 |
angles[1] = gaus_pl_pred - 90
|
47 |
-
angles[2] = gaus_ro_pred -
|
48 |
angles[3] = confidence
|
49 |
return angles
|
|
|
14 |
|
15 |
gaus_ax_pred = torch.argmax(dino_pred[:, 0:360], dim=-1)
|
16 |
gaus_pl_pred = torch.argmax(dino_pred[:, 360:360+180], dim=-1)
|
17 |
+
gaus_ro_pred = torch.argmax(dino_pred[:, 360+180:360+180+360], dim=-1)
|
18 |
confidence = F.softmax(dino_pred[:, -2:], dim=-1)[0][0]
|
19 |
angles = torch.zeros(4)
|
20 |
angles[0] = gaus_ax_pred
|
21 |
angles[1] = gaus_pl_pred - 90
|
22 |
+
angles[2] = gaus_ro_pred - 180
|
23 |
angles[3] = confidence
|
24 |
return angles
|
25 |
|
|
|
34 |
|
35 |
gaus_ax_pred = torch.argmax(dino_pred[:, 0:360], dim=-1).to(torch.float32)
|
36 |
gaus_pl_pred = torch.argmax(dino_pred[:, 360:360+180], dim=-1).to(torch.float32)
|
37 |
+
gaus_ro_pred = torch.argmax(dino_pred[:, 360+180:360+180+360], dim=-1).to(torch.float32)
|
38 |
|
39 |
gaus_ax_pred = remove_outliers_and_average_circular(gaus_ax_pred)
|
40 |
gaus_pl_pred = remove_outliers_and_average(gaus_pl_pred)
|
|
|
44 |
angles = torch.zeros(4)
|
45 |
angles[0] = gaus_ax_pred
|
46 |
angles[1] = gaus_pl_pred - 90
|
47 |
+
angles[2] = gaus_ro_pred - 180
|
48 |
angles[3] = confidence
|
49 |
return angles
|