Upload folder using huggingface_hub
Browse files
external/midas/__init__.py
CHANGED
@@ -1,20 +1,23 @@
|
|
1 |
import cv2
|
2 |
import numpy as np
|
3 |
import torch
|
4 |
-
|
5 |
from einops import rearrange
|
|
|
6 |
from .api import MiDaSInference
|
7 |
|
8 |
-
model =
|
9 |
|
10 |
|
11 |
def apply_midas(input_image, a=np.pi * 2.0, bg_th=0.1):
|
|
|
|
|
|
|
12 |
assert input_image.ndim == 3
|
13 |
image_depth = input_image
|
14 |
with torch.no_grad():
|
15 |
image_depth = torch.from_numpy(image_depth).float().cuda()
|
16 |
image_depth = image_depth / 127.5 - 1.0
|
17 |
-
image_depth = rearrange(image_depth,
|
18 |
depth = model(image_depth)[0]
|
19 |
|
20 |
depth_pt = depth.clone()
|
@@ -30,7 +33,7 @@ def apply_midas(input_image, a=np.pi * 2.0, bg_th=0.1):
|
|
30 |
x[depth_pt < bg_th] = 0
|
31 |
y[depth_pt < bg_th] = 0
|
32 |
normal = np.stack([x, y, z], axis=2)
|
33 |
-
normal /= np.sum(normal
|
34 |
normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
|
35 |
|
36 |
return depth_image, normal_image
|
|
|
1 |
import cv2
|
2 |
import numpy as np
|
3 |
import torch
|
|
|
4 |
from einops import rearrange
|
5 |
+
|
6 |
from .api import MiDaSInference
|
7 |
|
8 |
+
model = None
|
9 |
|
10 |
|
11 |
def apply_midas(input_image, a=np.pi * 2.0, bg_th=0.1):
|
12 |
+
global model
|
13 |
+
if not model:
|
14 |
+
model = MiDaSInference(model_type="dpt_hybrid").cuda()
|
15 |
assert input_image.ndim == 3
|
16 |
image_depth = input_image
|
17 |
with torch.no_grad():
|
18 |
image_depth = torch.from_numpy(image_depth).float().cuda()
|
19 |
image_depth = image_depth / 127.5 - 1.0
|
20 |
+
image_depth = rearrange(image_depth, "h w c -> 1 c h w")
|
21 |
depth = model(image_depth)[0]
|
22 |
|
23 |
depth_pt = depth.clone()
|
|
|
33 |
x[depth_pt < bg_th] = 0
|
34 |
y[depth_pt < bg_th] = 0
|
35 |
normal = np.stack([x, y, z], axis=2)
|
36 |
+
normal /= np.sum(normal**2.0, axis=2, keepdims=True) ** 0.5
|
37 |
normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
|
38 |
|
39 |
return depth_image, normal_image
|
internals/pipelines/replace_background.py
CHANGED
@@ -172,14 +172,5 @@ class ReplaceBackground(AbstractPipeline):
|
|
172 |
if not has_nsfw:
|
173 |
for i in range(len(images)):
|
174 |
images[i].paste(image, (0, 0), image)
|
175 |
-
w, h = images[i].size
|
176 |
-
out_bytes = self.upscaler.upscale(
|
177 |
-
image=images[i],
|
178 |
-
width=w,
|
179 |
-
height=h,
|
180 |
-
face_enhance=False,
|
181 |
-
resize_dimension=resize_dimension,
|
182 |
-
)
|
183 |
-
images[i] = Image.open(BytesIO(out_bytes)).convert("RGB")
|
184 |
|
185 |
return (images, has_nsfw)
|
|
|
172 |
if not has_nsfw:
|
173 |
for i in range(len(images)):
|
174 |
images[i].paste(image, (0, 0), image)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
175 |
|
176 |
return (images, has_nsfw)
|
internals/pipelines/upscaler.py
CHANGED
@@ -127,7 +127,7 @@ class Upscaler:
|
|
127 |
image = download_image(image)
|
128 |
|
129 |
w, h = image.size
|
130 |
-
if max(w, h) >
|
131 |
image = ImageUtil.resize_image(image, dimension=1024)
|
132 |
|
133 |
in_path = str(Path.home() / ".cache" / "input_upscale.png")
|
|
|
127 |
image = download_image(image)
|
128 |
|
129 |
w, h = image.size
|
130 |
+
if max(w, h) > 1024:
|
131 |
image = ImageUtil.resize_image(image, dimension=1024)
|
132 |
|
133 |
in_path = str(Path.home() / ".cache" / "input_upscale.png")
|