kimhyunwoo commited on
Commit
9106952
ยท
verified ยท
1 Parent(s): ee4e3d9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +56 -24
app.py CHANGED
@@ -4,7 +4,7 @@ from PIL import Image
4
  import numpy as np
5
  import torch
6
  from transformers import pipeline
7
- import cv2
8
 
9
  app = FastAPI()
10
 
@@ -12,48 +12,59 @@ app = FastAPI()
12
  # ๋ชจ๋ธ ๋กœ๋”ฉ์€ ์•ฑ ์‹œ์ž‘ ์‹œ ํ•œ ๋ฒˆ๋งŒ ํ•˜๋„๋ก ๊ธ€๋กœ๋ฒŒ ๋ณ€์ˆ˜๋กœ ์„ค์ •
13
  print("Loading Depth Anything model...")
14
  try:
15
- depth_estimator = pipeline(task="depth-estimation", model="LiangNX/depth-anything-hf")
 
 
16
  print("Depth Anything model loaded successfully.")
17
  except Exception as e:
18
  print(f"Error loading Depth Anything model: {e}")
19
  depth_estimator = None # ๋ชจ๋ธ ๋กœ๋“œ ์‹คํŒจ ์‹œ None์œผ๋กœ ์„ค์ •
20
 
21
  def process_image_for_depth(image_path_or_pil_image):
 
 
 
22
  if depth_estimator is None:
23
- return None, "Error: Depth Anything model not loaded."
24
 
25
  # Gradio๋Š” PIL Image ๊ฐ์ฒด๋กœ ์ด๋ฏธ์ง€๋ฅผ ์ „๋‹ฌํ•ฉ๋‹ˆ๋‹ค.
26
  if isinstance(image_path_or_pil_image, str):
 
27
  image = Image.open(image_path_or_pil_image).convert("RGB")
28
  else:
 
29
  image = image_path_or_pil_image.convert("RGB")
30
 
31
  try:
32
  # Depth Anything ๋ชจ๋ธ ์ถ”๋ก 
33
- # result๋Š” ๋”•์…”๋„ˆ๋ฆฌ๋กœ, 'depth' (PIL Image)์™€ 'depth_npy' (numpy array)๋ฅผ ํฌํ•จ
34
  result = depth_estimator(image)
35
 
36
- # ๋ށ์Šค ๋งต (PIL Image)
37
  depth_image_pil = result["depth"]
38
 
39
- # ๋ށ์Šค ๋งต (Numpy Array) - ์‹œ๊ฐํ™”๋ฅผ ์œ„ํ•ด ์ •๊ทœํ™”
40
- depth_np = result["depth_npy"]
41
- normalized_depth_np = (depth_np - depth_np.min()) / (depth_np.max() - depth_np.min()) * 255
42
- normalized_depth_np = normalized_depth_np.astype(np.uint8)
43
-
44
- # ํ‘๋ฐฑ ์ด๋ฏธ์ง€๋กœ ๋ณ€ํ™˜ (PIL Image)
45
- depth_grayscale_pil = Image.fromarray(normalized_depth_np)
46
 
47
- return depth_grayscale_pil, None
48
  except Exception as e:
 
 
49
  return None, f"Error processing image for depth: {e}"
50
 
51
  # Gradio ์ธํ„ฐํŽ˜์ด์Šค ์ •์˜
52
  with gr.Blocks() as demo:
53
  gr.Markdown("# ๐Ÿง‘โ€๐Ÿ’ป ์–ผ๊ตด ๋ށ์Šค ๋งต ์ถ”์ถœ๊ธฐ")
54
  gr.Markdown("์—ฌ๋Ÿฌ ์žฅ์˜ ์–ผ๊ตด ์‚ฌ์ง„์„ ์—…๋กœ๋“œํ•˜๋ฉด ๊ฐ ์‚ฌ์ง„์—์„œ ๋”ฅ๋Ÿฌ๋‹์„ ํ†ตํ•ด ๋ށ์Šค ๋งต(๊นŠ์ด ์ •๋ณด)์„ ์ถ”์ถœํ•ฉ๋‹ˆ๋‹ค.")
 
55
 
56
  with gr.Row():
 
 
57
  input_images = gr.File(label="์–ผ๊ตด ์‚ฌ์ง„ ์—…๋กœ๋“œ (์ตœ๋Œ€ 10์žฅ ๊ถŒ์žฅ)", file_count="multiple", type="filepath")
58
 
59
  output_gallery = gr.Gallery(label="์›๋ณธ ์ด๋ฏธ์ง€ ๋ฐ ๋ށ์Šค ๋งต", columns=[2], rows=[1], object_fit="contain", height="auto")
@@ -61,22 +72,41 @@ with gr.Blocks() as demo:
61
  process_button = gr.Button("๋ށ์Šค ๋งต ์ถ”์ถœ ์‹œ์ž‘")
62
 
63
  def process_all_images(image_paths):
 
 
 
64
  if not image_paths:
65
  return [(None, "์ด๋ฏธ์ง€๋ฅผ ์—…๋กœ๋“œํ•ด์ฃผ์„ธ์š”.")]
66
 
67
- results = []
 
68
  for i, path in enumerate(image_paths):
69
- original_image = Image.open(path).convert("RGB")
70
- depth_map_pil, error = process_image_for_depth(original_image)
71
-
72
- if error:
73
- print(f"Error processing image {i+1}: {error}")
74
- results.append((original_image, f"Error: {error}"))
75
- else:
76
- results.append((original_image, f"์›๋ณธ ์ด๋ฏธ์ง€ {i+1}"))
77
- results.append((depth_map_pil, f"๋ށ์Šค ๋งต {i+1}"))
78
- return results
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79
 
 
80
  process_button.click(
81
  fn=process_all_images,
82
  inputs=input_images,
@@ -84,9 +114,11 @@ with gr.Blocks() as demo:
84
  )
85
 
86
  # Gradio ์•ฑ์„ FastAPI์— ๋งˆ์šดํŠธ
 
87
  app = gr.mount_gradio_app(app, demo, path="/")
88
 
89
  # FastAPI ๊ธฐ๋ณธ ์—”๋“œํฌ์ธํŠธ (์„ ํƒ ์‚ฌํ•ญ, Gradio ์•ฑ์ด ๊ธฐ๋ณธ ๊ฒฝ๋กœ๋ฅผ ์ ์œ ํ•จ)
 
90
  @app.get("/api")
91
  def read_root():
92
  return {"message": "Welcome to the Face Depth Map Extractor! Visit / for the UI."}
 
4
  import numpy as np
5
  import torch
6
  from transformers import pipeline
7
+ import cv2 # cv2๋Š” ์ด๋ฏธ์ง€ ๋กœ๋“œ/์ €์žฅ/์ „์ฒ˜๋ฆฌ ๋“ฑ์— ์‚ฌ์šฉ๋  ์ˆ˜ ์žˆ์ง€๋งŒ, ํ˜„์žฌ ์ฝ”๋“œ์—์„œ๋Š” PIL๋งŒ์œผ๋กœ๋„ ์ถฉ๋ถ„ํ•ฉ๋‹ˆ๋‹ค.
8
 
9
  app = FastAPI()
10
 
 
12
  # ๋ชจ๋ธ ๋กœ๋”ฉ์€ ์•ฑ ์‹œ์ž‘ ์‹œ ํ•œ ๋ฒˆ๋งŒ ํ•˜๋„๋ก ๊ธ€๋กœ๋ฒŒ ๋ณ€์ˆ˜๋กœ ์„ค์ •
13
  print("Loading Depth Anything model...")
14
  try:
15
+ # ๐ŸŒŸ๐ŸŒŸ๐ŸŒŸ ๋ชจ๋ธ ์ด๋ฆ„ ์ˆ˜์ •: ์ •ํ™•ํ•œ Depth Anything v2 Large ๋ชจ๋ธ ID ๐ŸŒŸ๐ŸŒŸ๐ŸŒŸ
16
+ # ๋” ์ž‘์€ ๋ชจ๋ธ์„ ์›ํ•˜์‹œ๋ฉด "LiangNX/depth-anything-v2-base-nyu" ๋กœ ๋ณ€๊ฒฝํ•˜์„ธ์š”.
17
+ depth_estimator = pipeline(task="depth-estimation", model="LiangNX/depth-anything-v2-large-nyu", device="cpu") # GPU ์‚ฌ์šฉ ์‹œ device="cuda"
18
  print("Depth Anything model loaded successfully.")
19
  except Exception as e:
20
  print(f"Error loading Depth Anything model: {e}")
21
  depth_estimator = None # ๋ชจ๋ธ ๋กœ๋“œ ์‹คํŒจ ์‹œ None์œผ๋กœ ์„ค์ •
22
 
23
  def process_image_for_depth(image_path_or_pil_image):
24
+ """
25
+ ๋‹จ์ผ ์ด๋ฏธ์ง€์—์„œ ๋ށ์Šค ๋งต์„ ์ถ”์ถœํ•˜๋Š” ํ•จ์ˆ˜.
26
+ """
27
  if depth_estimator is None:
28
+ return None, "Error: Depth Anything model not loaded. Check server logs."
29
 
30
  # Gradio๋Š” PIL Image ๊ฐ์ฒด๋กœ ์ด๋ฏธ์ง€๋ฅผ ์ „๋‹ฌํ•ฉ๋‹ˆ๋‹ค.
31
  if isinstance(image_path_or_pil_image, str):
32
+ # ํŒŒ์ผ ๊ฒฝ๋กœ๋กœ ์ด๋ฏธ์ง€๊ฐ€ ์ „๋‹ฌ๋œ ๊ฒฝ์šฐ (์˜ˆ: Gradio์˜ `filepath` ํƒ€์ž…)
33
  image = Image.open(image_path_or_pil_image).convert("RGB")
34
  else:
35
+ # PIL Image ๊ฐ์ฒด๊ฐ€ ์ง์ ‘ ์ „๋‹ฌ๋œ ๊ฒฝ์šฐ
36
  image = image_path_or_pil_image.convert("RGB")
37
 
38
  try:
39
  # Depth Anything ๋ชจ๋ธ ์ถ”๋ก 
40
+ # result๋Š” ๋”•์…”๋„ˆ๋ฆฌ๋กœ, 'depth' (PIL Image)์™€ 'depth_npy' (numpy array)๋ฅผ ํฌํ•จํ•ฉ๋‹ˆ๋‹ค.
41
  result = depth_estimator(image)
42
 
43
+ # ๋ށ์Šค ๋งต (PIL Image) - ์ด PIL ์ด๋ฏธ์ง€๋Š” ์ด๋ฏธ ์‹œ๊ฐํ™”ํ•˜๊ธฐ ์ข‹์€ ํ˜•ํƒœ๋กœ ๋˜์–ด ์žˆ์Šต๋‹ˆ๋‹ค.
44
  depth_image_pil = result["depth"]
45
 
46
+ # ๋ށ์Šค ๋งต (Numpy Array) - ํ•„์š”ํ•˜๋‹ค๋ฉด ์ถ”๊ฐ€ ์ฒ˜๋ฆฌ (์˜ˆ: ์ •๊ทœํ™”, ๋‹ค๋ฅธ ํ˜•์‹์œผ๋กœ ๋ณ€ํ™˜)
47
+ # ํ˜„์žฌ๋Š” ํ‘๋ฐฑ ์ด๋ฏธ์ง€๋กœ ๋ฐ”๋กœ ์‚ฌ์šฉํ•ด๋„ ์ถฉ๋ถ„ํ•˜๋ฏ€๋กœ, ์•„๋ž˜ ์ฝ”๋“œ๋Š” ์„ ํƒ ์‚ฌํ•ญ์ž…๋‹ˆ๋‹ค.
48
+ # depth_np = result["depth_npy"]
49
+ # normalized_depth_np = (depth_np - depth_np.min()) / (depth_np.max() - depth_np.min()) * 255
50
+ # normalized_depth_np = normalized_depth_np.astype(np.uint8)
51
+ # depth_grayscale_pil = Image.fromarray(normalized_depth_np)
 
52
 
53
+ return depth_image_pil, None # ๋ށ์Šค ๋งต PIL Image ๋ฐ˜ํ™˜
54
  except Exception as e:
55
+ import traceback
56
+ traceback.print_exc() # ์—๋Ÿฌ ๋ฐœ์ƒ ์‹œ ์ƒ์„ธ ํŠธ๋ ˆ์ด์Šค๋ฐฑ ์ถœ๋ ฅ
57
  return None, f"Error processing image for depth: {e}"
58
 
59
  # Gradio ์ธํ„ฐํŽ˜์ด์Šค ์ •์˜
60
  with gr.Blocks() as demo:
61
  gr.Markdown("# ๐Ÿง‘โ€๐Ÿ’ป ์–ผ๊ตด ๋ށ์Šค ๋งต ์ถ”์ถœ๊ธฐ")
62
  gr.Markdown("์—ฌ๋Ÿฌ ์žฅ์˜ ์–ผ๊ตด ์‚ฌ์ง„์„ ์—…๋กœ๋“œํ•˜๋ฉด ๊ฐ ์‚ฌ์ง„์—์„œ ๋”ฅ๋Ÿฌ๋‹์„ ํ†ตํ•ด ๋ށ์Šค ๋งต(๊นŠ์ด ์ •๋ณด)์„ ์ถ”์ถœํ•ฉ๋‹ˆ๋‹ค.")
63
+ gr.Markdown("โš ๏ธ **์ฐธ๊ณ :** ๋ฌด๋ฃŒ CPU ํ™˜๊ฒฝ์—์„œ๋Š” ๋ชจ๋ธ ๋กœ๋”ฉ ๋ฐ ์ด๋ฏธ์ง€ ์ฒ˜๋ฆฌ ์‹œ๊ฐ„์ด ๋งค์šฐ ์˜ค๋ž˜ ๊ฑธ๋ฆด ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค.")
64
 
65
  with gr.Row():
66
+ # file_count="multiple"๋กœ ์—ฌ๋Ÿฌ ํŒŒ์ผ ์—…๋กœ๋“œ ๊ฐ€๋Šฅํ•˜๊ฒŒ ์„ค์ •
67
+ # type="filepath"๋กœ ์„ค์ •ํ•˜์—ฌ Gradio๊ฐ€ ํŒŒ์ผ์„ ์ž„์‹œ ๊ฒฝ๋กœ์— ์ €์žฅํ•˜๋„๋ก ํ•จ
68
  input_images = gr.File(label="์–ผ๊ตด ์‚ฌ์ง„ ์—…๋กœ๋“œ (์ตœ๋Œ€ 10์žฅ ๊ถŒ์žฅ)", file_count="multiple", type="filepath")
69
 
70
  output_gallery = gr.Gallery(label="์›๋ณธ ์ด๋ฏธ์ง€ ๋ฐ ๋ށ์Šค ๋งต", columns=[2], rows=[1], object_fit="contain", height="auto")
 
72
  process_button = gr.Button("๋ށ์Šค ๋งต ์ถ”์ถœ ์‹œ์ž‘")
73
 
74
  def process_all_images(image_paths):
75
+ """
76
+ ์—…๋กœ๋“œ๋œ ๋ชจ๋“  ์ด๋ฏธ์ง€์— ๋Œ€ํ•ด ๋ށ์Šค ๋งต ์ถ”์ถœ์„ ์ˆ˜ํ–‰ํ•˜๊ณ  ๊ฒฐ๊ณผ๋ฅผ ๋ฐ˜ํ™˜ํ•˜๋Š” ํ•จ์ˆ˜.
77
+ """
78
  if not image_paths:
79
  return [(None, "์ด๋ฏธ์ง€๋ฅผ ์—…๋กœ๋“œํ•ด์ฃผ์„ธ์š”.")]
80
 
81
+ results_for_gallery = [] # Gradio Gallery์— ํ‘œ์‹œํ•  (PIL Image, Label) ํŠœํ”Œ ๋ฆฌ์ŠคํŠธ
82
+
83
  for i, path in enumerate(image_paths):
84
+ try:
85
+ original_image = Image.open(path).convert("RGB")
86
+ depth_map_pil, error = process_image_for_depth(original_image)
87
+
88
+ # ์›๋ณธ ์ด๋ฏธ์ง€๋ฅผ ๋จผ์ € ์ถ”๊ฐ€
89
+ results_for_gallery.append((original_image, f"์›๋ณธ ์ด๋ฏธ์ง€ {i+1}"))
90
+
91
+ if error:
92
+ print(f"Error processing image {i+1}: {error}")
93
+ # ์˜ค๋ฅ˜๊ฐ€ ๋ฐœ์ƒํ•œ ๊ฒฝ์šฐ ๋ށ์Šค ๋งต ๋Œ€์‹  ์˜ค๋ฅ˜ ๋ฉ”์‹œ์ง€ ํ‘œ์‹œ
94
+ # ์ด ๋ถ€๋ถ„์€ Gradio Gallery๊ฐ€ (Image, Label)์„ ๊ธฐ๋Œ€ํ•˜๋ฏ€๋กœ,
95
+ # ์˜ค๋ฅ˜ ๋ฉ”์‹œ์ง€๋ฅผ ํ…์ŠคํŠธ ์ด๋ฏธ์ง€๋กœ ๋งŒ๋“ค๊ฑฐ๋‚˜, ์˜ค๋ฅ˜ ์ด๋ฏธ์ง€๋ฅผ ์‚ฌ์šฉํ•  ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค.
96
+ # ์—ฌ๊ธฐ์„œ๋Š” ๊ฐ„๋‹จํžˆ ๋นˆ ์ด๋ฏธ์ง€์™€ ํ•จ๊ป˜ ์˜ค๋ฅ˜ ๋ฉ”์‹œ์ง€๋ฅผ ํ‘œ์‹œํ•˜๋„๋ก ํ•ฉ๋‹ˆ๋‹ค.
97
+ dummy_error_image = Image.new('RGB', original_image.size, color = 'red')
98
+ results_for_gallery.append((dummy_error_image, f"์˜ค๋ฅ˜ ๋ฐœ์ƒ: {error}"))
99
+ else:
100
+ results_for_gallery.append((depth_map_pil, f"๋ށ์Šค ๋งต {i+1}"))
101
+ except Exception as e:
102
+ # ํŒŒ์ผ ์ž์ฒด๋ฅผ ์—ฌ๋Š” ๋ฐ ๋ฌธ์ œ๊ฐ€ ๋ฐœ์ƒํ•œ ๊ฒฝ์šฐ
103
+ print(f"Failed to open or process file {path}: {e}")
104
+ dummy_error_image = Image.new('RGB', (200, 200), color = 'red') # ์ž‘์€ ์—๋Ÿฌ ์ด๋ฏธ์ง€
105
+ results_for_gallery.append((dummy_error_image, f"ํŒŒ์ผ ์ฒ˜๋ฆฌ ์˜ค๋ฅ˜: {e}"))
106
+
107
+ return results_for_gallery
108
 
109
+ # ๋ฒ„ํŠผ ํด๋ฆญ ์‹œ ํ•จ์ˆ˜ ์—ฐ๊ฒฐ
110
  process_button.click(
111
  fn=process_all_images,
112
  inputs=input_images,
 
114
  )
115
 
116
  # Gradio ์•ฑ์„ FastAPI์— ๋งˆ์šดํŠธ
117
+ # ๊ธฐ๋ณธ ๊ฒฝ๋กœ (path="/")์— Gradio ์•ฑ์„ ๋งˆ์šดํŠธํ•˜์—ฌ ์›น ํŽ˜์ด์ง€ ์ ‘์† ์‹œ ๋ฐ”๋กœ UI๊ฐ€ ๋ณด์ด๋„๋ก ํ•ฉ๋‹ˆ๋‹ค.
118
  app = gr.mount_gradio_app(app, demo, path="/")
119
 
120
  # FastAPI ๊ธฐ๋ณธ ์—”๋“œํฌ์ธํŠธ (์„ ํƒ ์‚ฌํ•ญ, Gradio ์•ฑ์ด ๊ธฐ๋ณธ ๊ฒฝ๋กœ๋ฅผ ์ ์œ ํ•จ)
121
+ # ์ด ์—”๋“œํฌ์ธํŠธ๋Š” /api ๊ฒฝ๋กœ๋กœ ์ ‘์†ํ–ˆ์„ ๋•Œ๋งŒ ์ž‘๋™ํ•ฉ๋‹ˆ๋‹ค.
122
  @app.get("/api")
123
  def read_root():
124
  return {"message": "Welcome to the Face Depth Map Extractor! Visit / for the UI."}