|
- diff --git a/annotator/hed/__init__.py b/annotator/hed/__init__.py
- index 42d8dc6..1587035 100644
- --- a/annotator/hed/__init__.py
- +++ b/annotator/hed/__init__.py
- @@ -1,8 +1,12 @@
- +import pathlib
- +
- import numpy as np
- import cv2
- import torch
- from einops import rearrange
-
- +root_dir = pathlib.Path(__file__).parents[2]
- +
-
- class Network(torch.nn.Module):
- def __init__(self):
- @@ -64,7 +68,7 @@ class Network(torch.nn.Module):
- torch.nn.Sigmoid()
- )
-
- - self.load_state_dict({strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.load('./annotator/ckpts/network-bsds500.pth').items()})
- + self.load_state_dict({strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.load(f'{root_dir}/annotator/ckpts/network-bsds500.pth').items()})
- # end
-
- def forward(self, tenInput):
- diff --git a/annotator/midas/api.py b/annotator/midas/api.py
- index 9fa305e..d8594ea 100644
- --- a/annotator/midas/api.py
- +++ b/annotator/midas/api.py
- @@ -1,5 +1,7 @@
- # based on https://github.com/isl-org/MiDaS
-
- +import pathlib
- +
- import cv2
- import torch
- import torch.nn as nn
- @@ -10,10 +12,11 @@ from .midas.midas_net import MidasNet
- from .midas.midas_net_custom import MidasNet_small
- from .midas.transforms import Resize, NormalizeImage, PrepareForNet
-
- +root_dir = pathlib.Path(__file__).parents[2]
-
- ISL_PATHS = {
- - "dpt_large": "annotator/ckpts/dpt_large-midas-2f21e586.pt",
- - "dpt_hybrid": "annotator/ckpts/dpt_hybrid-midas-501f0c75.pt",
- + "dpt_large": f"{root_dir}/annotator/ckpts/dpt_large-midas-2f21e586.pt",
- + "dpt_hybrid": f"{root_dir}/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt",
- "midas_v21": "",
- "midas_v21_small": "",
- }
- diff --git a/annotator/mlsd/__init__.py b/annotator/mlsd/__init__.py
- index 75db717..f310fe6 100644
- --- a/annotator/mlsd/__init__.py
- +++ b/annotator/mlsd/__init__.py
- @@ -1,3 +1,5 @@
- +import pathlib
- +
- import cv2
- import numpy as np
- import torch
- @@ -8,8 +10,9 @@ from .models.mbv2_mlsd_tiny import MobileV2_MLSD_Tiny
- from .models.mbv2_mlsd_large import MobileV2_MLSD_Large
- from .utils import pred_lines
-
- +root_dir = pathlib.Path(__file__).parents[2]
-
- -model_path = './annotator/ckpts/mlsd_large_512_fp32.pth'
- +model_path = f'{root_dir}/annotator/ckpts/mlsd_large_512_fp32.pth'
- model = MobileV2_MLSD_Large()
- model.load_state_dict(torch.load(model_path), strict=True)
- model = model.cuda().eval()
- diff --git a/annotator/openpose/__init__.py b/annotator/openpose/__init__.py
- index 47d50a5..2369eed 100644
- --- a/annotator/openpose/__init__.py
- +++ b/annotator/openpose/__init__.py
- @@ -1,4 +1,5 @@
- import os
- +import pathlib
- os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
-
- import torch
- @@ -7,8 +8,10 @@ from . import util
- from .body import Body
- from .hand import Hand
-
- -body_estimation = Body('./annotator/ckpts/body_pose_model.pth')
- -hand_estimation = Hand('./annotator/ckpts/hand_pose_model.pth')
- +root_dir = pathlib.Path(__file__).parents[2]
- +
- +body_estimation = Body(f'{root_dir}/annotator/ckpts/body_pose_model.pth')
- +hand_estimation = Hand(f'{root_dir}/annotator/ckpts/hand_pose_model.pth')
-
-
- def apply_openpose(oriImg, hand=False):
- diff --git a/annotator/uniformer/__init__.py b/annotator/uniformer/__init__.py
- index 500e53c..4061dbe 100644
- --- a/annotator/uniformer/__init__.py
- +++ b/annotator/uniformer/__init__.py
- @@ -1,9 +1,12 @@
- +import pathlib
- +
- from annotator.uniformer.mmseg.apis import init_segmentor, inference_segmentor, show_result_pyplot
- from annotator.uniformer.mmseg.core.evaluation import get_palette
-
- +root_dir = pathlib.Path(__file__).parents[2]
-
- -checkpoint_file = "annotator/ckpts/upernet_global_small.pth"
- -config_file = 'annotator/uniformer/exp/upernet_global_small/config.py'
- +checkpoint_file = f"{root_dir}/annotator/ckpts/upernet_global_small.pth"
- +config_file = f'{root_dir}/annotator/uniformer/exp/upernet_global_small/config.py'
- model = init_segmentor(config_file, checkpoint_file).cuda()
-
-
|