imduckman commited on
Commit
b7ccabd
1 Parent(s): 7807e94

update pipeline

Browse files
models/depth_normal_pipeline_clip.py CHANGED
@@ -26,6 +26,32 @@ from utils.normal_ensemble import ensemble_normals
26
  from utils.batch_size import find_batch_size
27
  import cv2
28
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
  class DepthNormalPipelineOutput(BaseOutput):
30
  """
31
  Output class for Marigold monocular depth prediction pipeline.
@@ -179,10 +205,11 @@ class DepthNormalEstimationPipeline(DiffusionPipeline):
179
 
180
  # Resize back to original resolution
181
  if match_input_res:
 
182
  pred_img = Image.fromarray(depth_pred)
183
- pred_img = pred_img.resize(input_size)
184
  depth_pred = np.asarray(pred_img)
185
- normal_pred = cv2.resize(chw2hwc(normal_pred), input_size, interpolation = cv2.INTER_NEAREST)
186
 
187
  # Clip output range: current size is the original size
188
  depth_pred = depth_pred.clip(0, 1)
 
26
  from utils.batch_size import find_batch_size
27
  import cv2
28
 
29
+ def get_save_width_height(input_image_size, depth_pred):
30
+ orig_width, orig_height = input_image_size
31
+ depth_height, depth_width = depth_pred.shape[:2]
32
+
33
+ print("Inside pipe!")
34
+ print("Depth size: {}x{}".format(depth_width, depth_height))
35
+
36
+ # Use height
37
+ ratio = orig_height / depth_height
38
+ new_width = depth_width * ratio
39
+ if new_width <= orig_width:
40
+ save_width = int(new_width)
41
+ save_height = int(orig_height)
42
+ else:
43
+ ratio = orig_width / depth_width
44
+ new_height = depth_height * ratio
45
+ if new_height <= orig_height:
46
+ save_width = int(orig_width)
47
+ save_height = int(new_height)
48
+ else:
49
+ print("WTF IS THIS??")
50
+ return None, None
51
+
52
+ print("New depth size: {}x{}".format(save_width, save_height))
53
+ return save_width, save_height
54
+
55
  class DepthNormalPipelineOutput(BaseOutput):
56
  """
57
  Output class for Marigold monocular depth prediction pipeline.
 
205
 
206
  # Resize back to original resolution
207
  if match_input_res:
208
+ save_width, save_height = get_save_width_height(input_size, depth_pred)
209
  pred_img = Image.fromarray(depth_pred)
210
+ pred_img = pred_img.resize((save_width, save_height))
211
  depth_pred = np.asarray(pred_img)
212
+ normal_pred = cv2.resize(chw2hwc(normal_pred), (save_width, save_height), interpolation = cv2.INTER_NEAREST)
213
 
214
  # Clip output range: current size is the original size
215
  depth_pred = depth_pred.clip(0, 1)