Yuliang commited on
Commit
4a4217c
1 Parent(s): cd9b314

update app.py

Browse files
.gitignore CHANGED
@@ -17,5 +17,4 @@ dist
17
  *egg-info
18
  *.so
19
  run.sh
20
- *.log
21
- gradio_cached_examples/
17
  *egg-info
18
  *.so
19
  run.sh
20
+ *.log
 
README.md CHANGED
@@ -1,6 +1,6 @@
1
  ---
2
  title: Unconstrained & Detailed Clothed Human Digitization (ECON + ControlNet)
3
- metaTitle: Avatarify from Photo
4
  emoji: 🤼
5
  colorFrom: green
6
  colorTo: pink
1
  ---
2
  title: Unconstrained & Detailed Clothed Human Digitization (ECON + ControlNet)
3
+ metaTitle: ECON-Avatarify from Photo
4
  emoji: 🤼
5
  colorFrom: green
6
  colorTo: pink
app.py CHANGED
@@ -6,21 +6,14 @@ import os
6
 
7
  import subprocess
8
 
9
- curr_dir = os.path.dirname(__file__)
10
-
11
  if os.getenv('SYSTEM') == 'spaces':
12
  # subprocess.run('pip install pyembree'.split())
13
  subprocess.run(
14
  'pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu116_pyt1130/download.html'
15
  .split()
16
  )
17
- subprocess.run(
18
- f"cd {curr_dir}/lib/common/libmesh && python setup.py build_ext --inplace".split()
19
- )
20
- subprocess.run(
21
- f"cd {curr_dir}/lib/common/libvoxelize && python setup.py build_ext --inplace".split()
22
- )
23
- subprocess.run(f"cd {curr_dir}".split())
24
 
25
  from apps.infer import generate_model, generate_video
26
 
@@ -134,6 +127,8 @@ async (image_in_img, prompt, image_file_live_opt, live_conditioning) => {
134
  # Constants
135
  low_threshold = 100
136
  high_threshold = 200
 
 
137
 
138
  # Models
139
  pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
@@ -162,8 +157,8 @@ hint_prompts = '''
162
  <strong>Hints</strong>: <br>
163
  best quality, extremely detailed, solid color background,
164
  super detail, high detail, edge lighting, soft focus,
165
- light and dark contrast, 8k, high detail, edge lighting,
166
- 3d, c4d, blender, oc renderer, ultra high definition, 3d rendering
167
  '''
168
 
169
 
@@ -213,8 +208,6 @@ def toggle(choice):
213
  examples_pose = glob.glob('examples/pose/*')
214
  examples_cloth = glob.glob('examples/cloth/*')
215
 
216
- default_step = 50
217
-
218
  with gr.Blocks() as demo:
219
  gr.Markdown(description)
220
 
@@ -255,7 +248,7 @@ with gr.Blocks() as demo:
255
  gallery_cache = gr.State()
256
  inp = gr.Image(type="filepath", label="Input Image for ECON")
257
  fitting_step = gr.inputs.Slider(
258
- 10, 100, step=10, label='Fitting steps', default=default_step
259
  )
260
 
261
  with gr.Row():
@@ -283,34 +276,34 @@ with gr.Blocks() as demo:
283
  gr.Examples(
284
  examples=list(examples_pose),
285
  inputs=[inp],
286
- cache_examples=False,
287
  fn=generate_model,
288
  outputs=out_lst,
289
- label="Hard Pose Exampels"
290
  )
 
291
  gr.Examples(
292
  examples=list(examples_cloth),
293
  inputs=[inp],
294
- cache_examples=False,
295
  fn=generate_model,
296
  outputs=out_lst,
297
- label="Loose Cloth Exampels"
298
  )
 
 
299
 
300
  with gr.Column():
301
- overlap_inp = gr.Image(type="filepath", label="Image Normal Overlap")
302
- with gr.Row():
303
- out_final = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Clothed human")
304
- out_smpl = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="SMPL-X body")
305
 
306
  out_final_obj = gr.State()
307
  vis_tensor_path = gr.State()
308
 
309
  with gr.Row():
310
  btn_video = gr.Button("Generate Video (~2min)")
311
- with gr.Row():
312
- out_vid = gr.Video(label="Shared on Twitter with #ECON")
313
-
314
  # with gr.Row():
315
  # btn_texture = gr.Button("Generate Full-texture")
316
 
@@ -345,12 +338,13 @@ with gr.Blocks() as demo:
345
  )
346
 
347
  btn_submit.click(fn=generate_model, inputs=[inp, fitting_step], outputs=out_lst)
 
348
  # btn_texture.click(
349
  # fn=generate_texture,
350
  # inputs=[out_final_obj, prompt, seed, guidance_scale],
351
  # outputs=[viewpoint_images, result_video, output_file, progress_text]
352
  # )
353
-
354
  demo.load(None, None, None, _js=load_js)
355
 
356
  if __name__ == "__main__":
@@ -359,4 +353,5 @@ if __name__ == "__main__":
359
  # auth=(os.environ['USER'], os.environ['PASSWORD']),
360
  # auth_message="Register at icon.is.tue.mpg.de to get HuggingFace username and password.")
361
 
 
362
  demo.launch(debug=True, enable_queue=True)
6
 
7
  import subprocess
8
 
 
 
9
  if os.getenv('SYSTEM') == 'spaces':
10
  # subprocess.run('pip install pyembree'.split())
11
  subprocess.run(
12
  'pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu116_pyt1130/download.html'
13
  .split()
14
  )
15
+ subprocess.run("python setup.py build_ext --inplace".split(), cwd="./lib/common/libmesh/")
16
+ subprocess.run("python setup.py build_ext --inplace".split(), cwd="./lib/common/libvoxelize/")
 
 
 
 
 
17
 
18
  from apps.infer import generate_model, generate_video
19
 
127
  # Constants
128
  low_threshold = 100
129
  high_threshold = 200
130
+ default_step = 50
131
+ cached = False
132
 
133
  # Models
134
  pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
157
  <strong>Hints</strong>: <br>
158
  best quality, extremely detailed, solid color background,
159
  super detail, high detail, edge lighting, soft focus,
160
+ light and dark contrast, 8k, edge lighting, 3d, c4d,
161
+ blender, oc renderer, ultra high definition, 3d rendering
162
  '''
163
 
164
 
208
  examples_pose = glob.glob('examples/pose/*')
209
  examples_cloth = glob.glob('examples/cloth/*')
210
 
 
 
211
  with gr.Blocks() as demo:
212
  gr.Markdown(description)
213
 
248
  gallery_cache = gr.State()
249
  inp = gr.Image(type="filepath", label="Input Image for ECON")
250
  fitting_step = gr.inputs.Slider(
251
+ 10, 100, step=10, label='Fitting steps (Slower yet Better-aligned SMPL-X)', default=default_step
252
  )
253
 
254
  with gr.Row():
276
  gr.Examples(
277
  examples=list(examples_pose),
278
  inputs=[inp],
279
+ cache_examples=cached,
280
  fn=generate_model,
281
  outputs=out_lst,
282
+ label="Hard Pose Examples"
283
  )
284
+
285
  gr.Examples(
286
  examples=list(examples_cloth),
287
  inputs=[inp],
288
+ cache_examples=cached,
289
  fn=generate_model,
290
  outputs=out_lst,
291
+ label="Loose Cloth Examples"
292
  )
293
+
294
+ out_vid = gr.Video(label="Shared on Twitter with #ECON")
295
 
296
  with gr.Column():
297
+ overlap_inp = gr.Image(type="filepath", label="Image Normal Overlap").style(height=400)
298
+ out_final = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Clothed human", elem_id="avatar")
299
+ out_smpl = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="SMPL-X body", elem_id="avatar")
 
300
 
301
  out_final_obj = gr.State()
302
  vis_tensor_path = gr.State()
303
 
304
  with gr.Row():
305
  btn_video = gr.Button("Generate Video (~2min)")
306
+
 
 
307
  # with gr.Row():
308
  # btn_texture = gr.Button("Generate Full-texture")
309
 
338
  )
339
 
340
  btn_submit.click(fn=generate_model, inputs=[inp, fitting_step], outputs=out_lst)
341
+
342
  # btn_texture.click(
343
  # fn=generate_texture,
344
  # inputs=[out_final_obj, prompt, seed, guidance_scale],
345
  # outputs=[viewpoint_images, result_video, output_file, progress_text]
346
  # )
347
+
348
  demo.load(None, None, None, _js=load_js)
349
 
350
  if __name__ == "__main__":
353
  # auth=(os.environ['USER'], os.environ['PASSWORD']),
354
  # auth_message="Register at icon.is.tue.mpg.de to get HuggingFace username and password.")
355
 
356
+ demo.queue(concurrency_count=1)
357
  demo.launch(debug=True, enable_queue=True)
apps/infer.py CHANGED
@@ -28,6 +28,7 @@ import torch
28
  import torchvision
29
  import trimesh
30
  from pytorch3d.ops import SubdivideMeshes
 
31
  from termcolor import colored
32
  from tqdm.auto import tqdm
33
 
@@ -47,6 +48,7 @@ from lib.net.geometry import rot6d_to_rotmat, rotation_matrix_to_angle_axis
47
 
48
  torch.backends.cudnn.benchmark = True
49
 
 
50
  def generate_video(vis_tensor_path):
51
 
52
  in_tensor = torch.load(vis_tensor_path)
@@ -60,13 +62,14 @@ def generate_video(vis_tensor_path):
60
  # self-rotated video
61
  tmp_path = vis_tensor_path.replace("_in_tensor.pt", "_tmp.mp4")
62
  out_path = vis_tensor_path.replace("_in_tensor.pt", ".mp4")
63
-
64
  render.load_meshes(verts_lst, faces_lst)
65
  render.get_rendered_video_multi(in_tensor, tmp_path)
66
-
67
- os.system(f'ffmpeg -y -loglevel quiet -stats -i {tmp_path} -c:v libx264 {out_path}')
68
-
69
- return out_path, out_path
 
70
 
71
  def generate_model(in_path, fitting_step=50):
72
 
@@ -87,7 +90,12 @@ def generate_model(in_path, fitting_step=50):
87
 
88
  # load normal model
89
  normal_net = Normal.load_from_checkpoint(
90
- cfg=cfg, checkpoint_path=cfg.normal_path, map_location=device, strict=False
 
 
 
 
 
91
  )
92
  normal_net = normal_net.to(device)
93
  normal_net.netG.eval()
@@ -111,7 +119,12 @@ def generate_model(in_path, fitting_step=50):
111
  if cfg.bni.use_ifnet:
112
  # load IFGeo model
113
  ifnet = IFGeo.load_from_checkpoint(
114
- cfg=cfg, checkpoint_path=cfg.ifnet_path, map_location=device, strict=False
 
 
 
 
 
115
  )
116
  ifnet = ifnet.to(device)
117
  ifnet.netG.eval()
@@ -644,15 +657,13 @@ def generate_model(in_path, fitting_step=50):
644
  overlap_path = img_overlap_path
645
  vis_tensor_path = osp.join(out_dir, cfg.name, f"vid/{data['name']}_in_tensor.pt")
646
 
647
- # clean all the variables
648
- for element in dir():
649
- if 'path' not in element:
650
- del locals()[element]
651
 
652
- import gc
653
- gc.collect()
654
- torch.cuda.empty_cache()
655
 
656
- return [
657
- smpl_glb_path, refine_glb_path, refine_obj_path, overlap_path, vis_tensor_path
658
- ]
28
  import torchvision
29
  import trimesh
30
  from pytorch3d.ops import SubdivideMeshes
31
+ from huggingface_hub import hf_hub_download
32
  from termcolor import colored
33
  from tqdm.auto import tqdm
34
 
48
 
49
  torch.backends.cudnn.benchmark = True
50
 
51
+
52
  def generate_video(vis_tensor_path):
53
 
54
  in_tensor = torch.load(vis_tensor_path)
62
  # self-rotated video
63
  tmp_path = vis_tensor_path.replace("_in_tensor.pt", "_tmp.mp4")
64
  out_path = vis_tensor_path.replace("_in_tensor.pt", ".mp4")
65
+
66
  render.load_meshes(verts_lst, faces_lst)
67
  render.get_rendered_video_multi(in_tensor, tmp_path)
68
+
69
+ os.system(f"ffmpeg -y -loglevel quiet -stats -i {tmp_path} -vcodec libx264 {out_path}")
70
+
71
+ return out_path
72
+
73
 
74
  def generate_model(in_path, fitting_step=50):
75
 
90
 
91
  # load normal model
92
  normal_net = Normal.load_from_checkpoint(
93
+ cfg=cfg,
94
+ checkpoint_path=hf_hub_download(
95
+ repo_id="Yuliang/ICON", use_auth_token=os.environ["ICON"], filename=cfg.normal_path
96
+ ),
97
+ map_location=device,
98
+ strict=False
99
  )
100
  normal_net = normal_net.to(device)
101
  normal_net.netG.eval()
119
  if cfg.bni.use_ifnet:
120
  # load IFGeo model
121
  ifnet = IFGeo.load_from_checkpoint(
122
+ cfg=cfg,
123
+ checkpoint_path=hf_hub_download(
124
+ repo_id="Yuliang/ICON", use_auth_token=os.environ["ICON"], filename=cfg.ifnet_path
125
+ ),
126
+ map_location=device,
127
+ strict=False
128
  )
129
  ifnet = ifnet.to(device)
130
  ifnet.netG.eval()
657
  overlap_path = img_overlap_path
658
  vis_tensor_path = osp.join(out_dir, cfg.name, f"vid/{data['name']}_in_tensor.pt")
659
 
660
+ # # clean all the variables
661
+ # for element in dir():
662
+ # if 'path' not in element:
663
+ # del locals()[element]
664
 
665
+ # import gc
666
+ # gc.collect()
667
+ # torch.cuda.empty_cache()
668
 
669
+ return [smpl_glb_path, refine_glb_path, refine_obj_path, overlap_path, vis_tensor_path]
 
 
configs/econ.yaml CHANGED
@@ -1,7 +1,7 @@
1
  name: econ
2
  ckpt_dir: "./data/ckpt/"
3
- normal_path: "./data/ckpt/normal.ckpt"
4
- ifnet_path: "./data/ckpt/ifnet.ckpt"
5
  results_path: "./results"
6
 
7
  net:
1
  name: econ
2
  ckpt_dir: "./data/ckpt/"
3
+ normal_path: "normal.ckpt"
4
+ ifnet_path: "ifnet.ckpt"
5
  results_path: "./results"
6
 
7
  net:
gradio_cached_examples/13/log.csv ADDED
@@ -0,0 +1,2 @@
 
 
1
+ flag,username,timestamp
2
+ ,,2023-04-15 18:15:46.412679
gradio_cached_examples/25/log.csv ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
1
+ flag,username,timestamp
2
+ ,,2023-04-16 10:48:00.715491
3
+ ,,2023-04-16 10:50:02.250539
4
+ ,,2023-04-16 10:52:15.683112
5
+ ,,2023-04-16 10:54:18.253116
6
+ ,,2023-04-16 10:56:22.892765
lib/common/libmesh/inside_mesh.py CHANGED
@@ -1,5 +1,4 @@
1
  import numpy as np
2
-
3
  from .triangle_hash import TriangleHash as _TriangleHash
4
 
5
 
1
  import numpy as np
 
2
  from .triangle_hash import TriangleHash as _TriangleHash
3
 
4
 
lib/common/render.py CHANGED
@@ -38,6 +38,7 @@ from pytorch3d.renderer import (
38
  )
39
  from pytorch3d.renderer.mesh import TexturesVertex
40
  from pytorch3d.structures import Meshes
 
41
  from termcolor import colored
42
  from tqdm import tqdm
43
 
@@ -305,6 +306,9 @@ class Render:
305
 
306
  height, width = data["img_raw"].shape[2:]
307
 
 
 
 
308
  fourcc = cv2.VideoWriter_fourcc(*"mp4v")
309
  video = cv2.VideoWriter(
310
  save_path,
@@ -351,9 +355,12 @@ class Render:
351
  data)
352
  img_cloth = blend_rgb_norm((torch.stack(mesh_renders)[num_obj:, cam_id] - 0.5) * 2.0,
353
  data)
354
- final_img = torch.cat([img_raw, img_smpl, img_cloth],
355
- dim=-1).squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)
 
 
 
356
 
357
- video.write(final_img[:, :, ::-1])
358
 
359
  video.release()
38
  )
39
  from pytorch3d.renderer.mesh import TexturesVertex
40
  from pytorch3d.structures import Meshes
41
+ import torch.nn.functional as F
42
  from termcolor import colored
43
  from tqdm import tqdm
44
 
306
 
307
  height, width = data["img_raw"].shape[2:]
308
 
309
+ width = int(width / (height / 256.0))
310
+ height = 256
311
+
312
  fourcc = cv2.VideoWriter_fourcc(*"mp4v")
313
  video = cv2.VideoWriter(
314
  save_path,
355
  data)
356
  img_cloth = blend_rgb_norm((torch.stack(mesh_renders)[num_obj:, cam_id] - 0.5) * 2.0,
357
  data)
358
+ final_img = torch.cat([img_raw, img_smpl, img_cloth], dim=-1).squeeze(0)
359
+
360
+ final_img_rescale = F.interpolate(
361
+ final_img, size=(height, width), mode="bilinear", align_corners=False
362
+ ).squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)
363
 
364
+ video.write(final_img_rescale[:, :, ::-1])
365
 
366
  video.release()