update model weight

app.py

@@ -8,7 +8,7 @@ from inference import *
 from utils import *
 
 from huggingface_hub import hf_hub_download
-ckpt_path = hf_hub_download(repo_id="Viglong/Orient-Anything", filename="croplargeEX03/dino_weight.pt", repo_type="model", cache_dir='./', resume_download=True)
+ckpt_path = hf_hub_download(repo_id="Viglong/Orient-Anything", filename="ronormsigma1/dino_weight.pt", repo_type="model", cache_dir='./', resume_download=True)
 print(ckpt_path)
 
 save_path = './'
@@ -16,7 +16,7 @@ device = 'cpu'
 dino = DINOv2_MLP(
                     dino_mode   = 'large',
                     in_dim      = 1024,
-                    out_dim     = 360+180+180+2,
+                    out_dim     = 360+180+360+2,
                     evaluate    = True,
                     mask_dino   = False,
                     frozen_back = False

inference.py

@@ -14,12 +14,12 @@ def get_3angle(image, dino, val_preprocess, device):
 
     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)
+    gaus_ro_pred   = torch.argmax(dino_pred[:, 360+180:360+180+360], 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[2]  = gaus_ro_pred - 180
     angles[3]  = confidence
     return angles
 
@@ -34,7 +34,7 @@ def get_3angle_infer_aug(origin_img, rm_bkg_img, dino, val_preprocess, device):
 
     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_ro_pred   = torch.argmax(dino_pred[:, 360+180:360+180+360], 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)
@@ -44,6 +44,6 @@ def get_3angle_infer_aug(origin_img, rm_bkg_img, dino, val_preprocess, device):
     angles = torch.zeros(4)
     angles[0]  = gaus_ax_pred
     angles[1]  = gaus_pl_pred - 90
-    angles[2]  = gaus_ro_pred - 90
+    angles[2]  = gaus_ro_pred - 180
     angles[3]  = confidence
     return angles