Spaces:
Running
Running
File size: 4,358 Bytes
96b3fb1 e189cea |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
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()
diff --git a/annotator/util.py b/annotator/util.py
index 7cde937..10a6d58 100644
--- a/annotator/util.py
+++ b/annotator/util.py
@@ -25,7 +25,7 @@ def resize_image(input_image, resolution):
H, W, C = input_image.shape
H = float(H)
W = float(W)
- k = float(resolution) / min(H, W)
+ k = float(resolution) / max(H, W)
H *= k
W *= k
H = int(np.round(H / 64.0)) * 64
|