alexnasa commited on
Commit
c22a956
·
verified ·
1 Parent(s): a79ad4c

Delete app_photo.py

Browse files
Files changed (1) hide show
  1. app_photo.py +0 -224
app_photo.py DELETED
@@ -1,224 +0,0 @@
1
- import spaces
2
- import os
3
- import subprocess
4
- import tempfile
5
- import uuid
6
- import glob
7
- import shutil
8
- import time
9
- import gradio as gr
10
- import sys
11
- from PIL import Image
12
-
13
- def install_cuda_toolkit():
14
- CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
15
- CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
16
- subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
17
- subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
18
- subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
19
-
20
- os.environ["CUDA_HOME"] = "/usr/local/cuda"
21
- os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
22
- os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
23
- os.environ["CUDA_HOME"],
24
- "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
25
- )
26
- # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
27
- os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
28
- print("==> finished installation")
29
-
30
- install_cuda_toolkit()
31
-
32
-
33
- import os
34
- import torch
35
- import numpy as np
36
- import trimesh
37
- from pytorch3d.io import load_obj
38
- from pixel3dmm.tracking.renderer_nvdiffrast import NVDRenderer
39
- from pixel3dmm.tracking.flame.FLAME import FLAME
40
- from pixel3dmm.tracking.tracker import Tracker
41
- from pixel3dmm import env_paths
42
- from omegaconf import OmegaConf
43
-
44
-
45
- DEVICE = "cuda"
46
-
47
- base_conf = OmegaConf.load(f'{env_paths.CODE_BASE}/configs/tracking.yaml')
48
-
49
- _mesh_file = env_paths.head_template
50
- flame_model = FLAME(base_conf).to(DEVICE)
51
-
52
- _obj_faces = load_obj(_mesh_file)[1]
53
-
54
- diff_renderer = NVDRenderer(
55
- image_size=base_conf.size,
56
- obj_filename=_mesh_file,
57
- no_sh=False,
58
- white_bg=True
59
- ).to(DEVICE)
60
-
61
-
62
- # Utility to select first image from a folder
63
- def first_image_from_dir(directory):
64
- patterns = ["*.jpg", "*.png", "*.jpeg"]
65
- files = []
66
- for p in patterns:
67
- files.extend(glob.glob(os.path.join(directory, p)))
68
- if not files:
69
- return None
70
- return sorted(files)[0]
71
-
72
- # Function to reset the UI and state
73
- def reset_all():
74
- return (
75
- None, # crop_img
76
- None, # normals_img
77
- None, # uv_img
78
- None, # track_img
79
- "Awaiting new image upload...", # status
80
- {}, # state
81
- gr.update(interactive=True), # preprocess_btn
82
- gr.update(interactive=False), # normals_btn
83
- gr.update(interactive=False), # uv_map_btn
84
- gr.update(interactive=False) # track_btn
85
- )
86
-
87
- # Step 1: Preprocess the input image (Save and Crop)
88
- # @spaces.GPU()
89
- def preprocess_image(image_array, state):
90
- if image_array is None:
91
- return "❌ Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=False)
92
-
93
- session_id = str(uuid.uuid4())
94
- base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
95
- os.makedirs(base_dir, exist_ok=True)
96
- state.update({"session_id": session_id, "base_dir": base_dir})
97
-
98
- img = Image.fromarray(image_array)
99
- saved_image_path = os.path.join(base_dir, f"{session_id}.png")
100
- img.save(saved_image_path)
101
- state["image_path"] = saved_image_path
102
-
103
- try:
104
- p = subprocess.run([
105
- "python", "scripts/run_preprocessing.py", "--video_or_images_path", saved_image_path
106
- ], check=True, capture_output=True, text=True)
107
- except subprocess.CalledProcessError as e:
108
- err = f"❌ Preprocess failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
109
- shutil.rmtree(base_dir)
110
- return err, None, {}, gr.update(interactive=True), gr.update(interactive=False)
111
-
112
- crop_dir = os.path.join(base_dir, "cropped")
113
- image = first_image_from_dir(crop_dir)
114
- return "✅ Step 1 complete. Ready for Normals.", image, state, gr.update(interactive=False), gr.update(interactive=True)
115
-
116
- # Step 2: Normals inference → normals image
117
- @spaces.GPU()
118
- def step2_normals(state):
119
- session_id = state.get("session_id")
120
- if not session_id:
121
- return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False), gr.update(interactive=False)
122
-
123
- try:
124
- p = subprocess.run([
125
- "python", "scripts/network_inference.py", "model.prediction_type=normals", f"video_name={session_id}"
126
- ], check=True, capture_output=True, text=True)
127
- except subprocess.CalledProcessError as e:
128
- err = f"❌ Normal map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
129
- return err, None, state, gr.update(interactive=True), gr.update(interactive=False)
130
-
131
- normals_dir = os.path.join(state["base_dir"], "p3dmm", "normals")
132
- image = first_image_from_dir(normals_dir)
133
- return "✅ Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive=False), gr.update(interactive=True)
134
-
135
- # Step 3: UV map inference → uv map image
136
- @spaces.GPU()
137
- def step3_uv_map(state):
138
- session_id = state.get("session_id")
139
- if not session_id:
140
- return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False), gr.update(interactive=False)
141
-
142
- try:
143
- p = subprocess.run([
144
- "python", "scripts/network_inference.py", "model.prediction_type=uv_map", f"video_name={session_id}"
145
- ], check=True, capture_output=True, text=True)
146
- except subprocess.CalledProcessError as e:
147
- err = f"❌ UV map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
148
- return err, None, state, gr.update(interactive=True), gr.update(interactive=False)
149
-
150
- uv_dir = os.path.join(state["base_dir"], "p3dmm", "uv_map")
151
- image = first_image_from_dir(uv_dir)
152
- return "✅ Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive=False), gr.update(interactive=True)
153
-
154
- # Step 4: Tracking → final tracking image
155
- @spaces.GPU()
156
- def step4_track(state):
157
- session_id = state.get("session_id")
158
- base_conf.video_name = f'{session_id}'
159
- tracker = Tracker(base_conf, flame_model, diff_renderer)
160
- tracker.run()
161
-
162
- tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames")
163
- image = first_image_from_dir(tracking_dir)
164
-
165
- return "✅ Pipeline complete!", image, state, gr.update(interactive=False)
166
-
167
- # Build Gradio UI
168
- demo = gr.Blocks()
169
-
170
- with demo:
171
- gr.Markdown("## Image Processing Pipeline")
172
- gr.Markdown("Upload an image, then click the buttons in order. Uploading a new image will reset the process.")
173
- with gr.Row():
174
- with gr.Column():
175
- image_in = gr.Image(label="Upload Image", type="numpy", height=512)
176
- status = gr.Textbox(label="Status", lines=2, interactive=False, value="Upload an image to start.")
177
- state = gr.State({})
178
- with gr.Column():
179
- with gr.Row():
180
- crop_img = gr.Image(label="Preprocessed", height=256)
181
- normals_img = gr.Image(label="Normals", height=256)
182
- with gr.Row():
183
- uv_img = gr.Image(label="UV Map", height=256)
184
- track_img = gr.Image(label="Tracking", height=256)
185
-
186
- with gr.Row():
187
- preprocess_btn = gr.Button("Step 1: Preprocess", interactive=True)
188
- normals_btn = gr.Button("Step 2: Normals", interactive=False)
189
- uv_map_btn = gr.Button("Step 3: UV Map", interactive=False)
190
- track_btn = gr.Button("Step 4: Track", interactive=False)
191
-
192
- # Define component list for reset
193
- outputs_for_reset = [crop_img, normals_img, uv_img, track_img, status, state, preprocess_btn, normals_btn, uv_map_btn, track_btn]
194
-
195
- # Pipeline execution logic
196
- preprocess_btn.click(
197
- fn=preprocess_image,
198
- inputs=[image_in, state],
199
- outputs=[status, crop_img, state, preprocess_btn, normals_btn]
200
- )
201
- normals_btn.click(
202
- fn=step2_normals,
203
- inputs=[state],
204
- outputs=[status, normals_img, state, normals_btn, uv_map_btn]
205
- )
206
- uv_map_btn.click(
207
- fn=step3_uv_map,
208
- inputs=[state],
209
- outputs=[status, uv_img, state, uv_map_btn, track_btn]
210
- )
211
- track_btn.click(
212
- fn=step4_track,
213
- inputs=[state],
214
- outputs=[status, track_img, state, track_btn]
215
- )
216
-
217
- # Event to reset everything when a new image is uploaded
218
- image_in.upload(fn=reset_all, inputs=None, outputs=outputs_for_reset)
219
-
220
- # ------------------------------------------------------------------
221
- # START THE GRADIO SERVER
222
- # ------------------------------------------------------------------
223
- demo.queue()
224
- demo.launch(share=True, ssr_mode=False)