adpro commited on
Commit
02acfac
·
verified ·
1 Parent(s): bd0efcd

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +34 -23
app.py CHANGED
@@ -3,47 +3,58 @@ from transformers import DPTFeatureExtractor, DPTForDepthEstimation
3
  import torch
4
  import numpy as np
5
  from PIL import Image
 
6
  from pathlib import Path
 
7
 
8
- # Load model and feature extractor
9
  feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
10
  model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
11
- model.eval()
12
 
13
- def process_image(image):
14
- # Chuẩn hóa ảnh đầu vào
 
 
 
 
 
 
 
 
15
  encoding = feature_extractor(image, return_tensors="pt")
16
-
17
- # Forward qua model
18
  with torch.no_grad():
19
  outputs = model(**encoding)
20
  predicted_depth = outputs.predicted_depth
21
-
22
- # Resize output về đúng kích thước ảnh gốc
23
  prediction = torch.nn.functional.interpolate(
24
  predicted_depth.unsqueeze(1),
25
- size=image.size[::-1], # (H, W)
26
  mode="bicubic",
27
- align_corners=False
28
  ).squeeze()
29
-
30
- # Chuyển thành ảnh uint8
31
  output = prediction.cpu().numpy()
32
- formatted = (output * 255 / np.max(output)).astype('uint8')
33
- img = Image.fromarray(formatted)
34
  return [img]
35
 
36
- # Interface
37
- title = "Demo: Zero-shot Depth Estimation with DPT"
38
- description = "Intel's DPT: Dense Prediction Transformer for depth estimation from a single image."
 
 
39
 
40
  iface = gr.Interface(
41
- fn=process_image,
42
- inputs=gr.inputs.Image(type="pil", label="Input Image"),
43
- outputs=predicted_depth,
 
 
44
  title=title,
45
  description=description,
46
- allow_flagging="never"
 
 
47
  )
48
-
49
- iface.launch(debug=True)
 
3
  import torch
4
  import numpy as np
5
  from PIL import Image
6
+ import open3d as o3d
7
  from pathlib import Path
8
+ import os
9
 
 
10
  feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
11
  model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
 
12
 
13
+
14
+ def process_image(image_path):
15
+ image_path = Path(image_path)
16
+ image_raw = Image.open(image_path)
17
+ image = image_raw.resize(
18
+ (800, int(800 * image_raw.size[1] / image_raw.size[0])),
19
+ Image.Resampling.LANCZOS,
20
+ )
21
+
22
+ # prepare image for the model
23
  encoding = feature_extractor(image, return_tensors="pt")
24
+
25
+ # forward pass
26
  with torch.no_grad():
27
  outputs = model(**encoding)
28
  predicted_depth = outputs.predicted_depth
29
+
30
+ # interpolate to original size
31
  prediction = torch.nn.functional.interpolate(
32
  predicted_depth.unsqueeze(1),
33
+ size=image.size[::-1],
34
  mode="bicubic",
35
+ align_corners=False,
36
  ).squeeze()
 
 
37
  output = prediction.cpu().numpy()
38
+ depth_image = (output * 255 / np.max(output)).astype("uint8")
39
+ img = Image.fromarray(depth_image)
40
  return [img]
41
 
42
+
43
+
44
+ title = "Demo: zero-shot depth estimation with DPT + 3D Point Cloud"
45
+ description = "This demo is a variation from the original <a href='https://huggingface.co/spaces/nielsr/dpt-depth-estimation' target='_blank'>DPT Demo</a>. It uses the DPT model to predict the depth of an image and then uses 3D Point Cloud to create a 3D object."
46
+ examples = [["examples/" + img] for img in os.listdir("examples/")]
47
 
48
  iface = gr.Interface(
49
+ fn=process_image,
50
+ inputs=[gr.Image(type="filepath", label="Input Image")],
51
+ outputs=[
52
+ gr.Image(label="predicted depth", type="pil"),
53
+ ],
54
  title=title,
55
  description=description,
56
+ examples=examples,
57
+ allow_flagging="never",
58
+ cache_examples=False,
59
  )
60
+ iface.launch(debug=True, show_api=False)