alexnasa commited on
Commit
6067df7
Β·
verified Β·
1 Parent(s): ffb4d06

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -37
app.py CHANGED
@@ -15,35 +15,6 @@ os.environ["PIXEL3DMM_CODE_BASE"] = f"{os.getcwd()}"
15
  os.environ["PIXEL3DMM_PREPROCESSED_DATA"] = f"{os.getcwd()}/proprocess_results"
16
  os.environ["PIXEL3DMM_TRACKING_OUTPUT"] = f"{os.getcwd()}/tracking_results"
17
 
18
-
19
- def sh(cmd): subprocess.check_call(cmd, shell=True)
20
-
21
- # only do this once per VM restart
22
- sh("pip install -e .")
23
- sh("cd src/pixel3dmm/preprocessing/facer && pip install -e .")
24
- sh("cd src/pixel3dmm/preprocessing/PIPNet/FaceBoxesV2/utils && sh make.sh")
25
-
26
-
27
- def install_cuda_toolkit():
28
- CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
29
- CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
30
- subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
31
- subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
32
- subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
33
-
34
- os.environ["CUDA_HOME"] = "/usr/local/cuda"
35
- os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
36
- os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
37
- os.environ["CUDA_HOME"],
38
- "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
39
- )
40
- # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
41
- os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
42
- print("==> finished installation")
43
-
44
-
45
- install_cuda_toolkit()
46
-
47
  # Utility to select first image from a folder
48
  def first_image_from_dir(directory):
49
  patterns = ["*.jpg", "*.png", "*.jpeg"]
@@ -73,7 +44,7 @@ def reset_all():
73
  @spaces.GPU()
74
  def preprocess_image(image_array, state):
75
  if image_array is None:
76
- return "❌ Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=True)
77
 
78
  session_id = str(uuid.uuid4())
79
  base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
@@ -96,14 +67,14 @@ def preprocess_image(image_array, state):
96
 
97
  crop_dir = os.path.join(base_dir, "cropped")
98
  image = first_image_from_dir(crop_dir)
99
- return "βœ… Step 1 complete. Ready for Normals.", image, state, gr.update(interactive=True), gr.update(interactive=True)
100
 
101
  # Step 2: Normals inference β†’ normals image
102
  @spaces.GPU()
103
  def step2_normals(state):
104
  session_id = state.get("session_id")
105
  if not session_id:
106
- return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=True), gr.update(interactive=False)
107
 
108
  try:
109
  p = subprocess.run([
@@ -115,14 +86,14 @@ def step2_normals(state):
115
 
116
  normals_dir = os.path.join(state["base_dir"], "p3dmm", "normals")
117
  image = first_image_from_dir(normals_dir)
118
- return "βœ… Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive=True), gr.update(interactive=True)
119
 
120
  # Step 3: UV map inference β†’ uv map image
121
  @spaces.GPU()
122
  def step3_uv_map(state):
123
  session_id = state.get("session_id")
124
  if not session_id:
125
- return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=True), gr.update(interactive=True)
126
 
127
  try:
128
  p = subprocess.run([
@@ -134,14 +105,14 @@ def step3_uv_map(state):
134
 
135
  uv_dir = os.path.join(state["base_dir"], "p3dmm", "uv_map")
136
  image = first_image_from_dir(uv_dir)
137
- return "βœ… Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive=True), gr.update(interactive=True)
138
 
139
  # Step 4: Tracking β†’ final tracking image
140
  @spaces.GPU()
141
  def step4_track(state):
142
  session_id = state.get("session_id")
143
  if not session_id:
144
- return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=True)
145
 
146
  script = os.path.join(os.environ["PIXEL3DMM_CODE_BASE"], "scripts", "track.py")
147
  try:
@@ -154,7 +125,7 @@ def step4_track(state):
154
 
155
  tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames")
156
  image = first_image_from_dir(tracking_dir)
157
- return "βœ… Pipeline complete!", image, state, gr.update(interactive=True)
158
 
159
  # Build Gradio UI
160
  demo = gr.Blocks()
 
15
  os.environ["PIXEL3DMM_PREPROCESSED_DATA"] = f"{os.getcwd()}/proprocess_results"
16
  os.environ["PIXEL3DMM_TRACKING_OUTPUT"] = f"{os.getcwd()}/tracking_results"
17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
  # Utility to select first image from a folder
19
  def first_image_from_dir(directory):
20
  patterns = ["*.jpg", "*.png", "*.jpeg"]
 
44
  @spaces.GPU()
45
  def preprocess_image(image_array, state):
46
  if image_array is None:
47
+ return "❌ Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=False)
48
 
49
  session_id = str(uuid.uuid4())
50
  base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
 
67
 
68
  crop_dir = os.path.join(base_dir, "cropped")
69
  image = first_image_from_dir(crop_dir)
70
+ return "βœ… Step 1 complete. Ready for Normals.", image, state, gr.update(interactive=False), gr.update(interactive=True)
71
 
72
  # Step 2: Normals inference β†’ normals image
73
  @spaces.GPU()
74
  def step2_normals(state):
75
  session_id = state.get("session_id")
76
  if not session_id:
77
+ return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False), gr.update(interactive=False)
78
 
79
  try:
80
  p = subprocess.run([
 
86
 
87
  normals_dir = os.path.join(state["base_dir"], "p3dmm", "normals")
88
  image = first_image_from_dir(normals_dir)
89
+ return "βœ… Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive=False), gr.update(interactive=True)
90
 
91
  # Step 3: UV map inference β†’ uv map image
92
  @spaces.GPU()
93
  def step3_uv_map(state):
94
  session_id = state.get("session_id")
95
  if not session_id:
96
+ return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False), gr.update(interactive=False)
97
 
98
  try:
99
  p = subprocess.run([
 
105
 
106
  uv_dir = os.path.join(state["base_dir"], "p3dmm", "uv_map")
107
  image = first_image_from_dir(uv_dir)
108
+ return "βœ… Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive=False), gr.update(interactive=True)
109
 
110
  # Step 4: Tracking β†’ final tracking image
111
  @spaces.GPU()
112
  def step4_track(state):
113
  session_id = state.get("session_id")
114
  if not session_id:
115
+ return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False)
116
 
117
  script = os.path.join(os.environ["PIXEL3DMM_CODE_BASE"], "scripts", "track.py")
118
  try:
 
125
 
126
  tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames")
127
  image = first_image_from_dir(tracking_dir)
128
+ return "βœ… Pipeline complete!", image, state, gr.update(interactive=False)
129
 
130
  # Build Gradio UI
131
  demo = gr.Blocks()