jayparmr commited on
Commit
92207d3
1 Parent(s): 1377831

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 = MiDaSInference(model_type="dpt_hybrid").cuda()
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, 'h w c -> 1 c h w')
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 ** 2.0, axis=2, keepdims=True) ** 0.5
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) > 1536:
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")