299 lines
10 KiB
Python
299 lines
10 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
import platform
|
|
import pyrender
|
|
# from pathlib import Path
|
|
from tqdm import tqdm
|
|
|
|
|
|
if "PYOPENGL_PLATFORM" not in os.environ:
|
|
os.environ["PYOPENGL_PLATFORM"] = "egl"
|
|
|
|
import comfy.utils
|
|
|
|
from ..humans4d.hmr2.models import load_hmr2
|
|
from ..humans4d.hmr2.utils import recursive_to
|
|
from ..humans4d.hmr2.datasets.vitdet_dataset import ViTDetDataset
|
|
from ..humans4d.hmr2.utils.renderer import cam_crop_to_full
|
|
# from ..humans4d.hmr2.configs import get_config
|
|
# from ..humans4d.hmr2.models import check_smpl_exists
|
|
from ..humans4d.hmr2.utils.utils_detectron2 import DefaultPredictor_Lazy
|
|
from detectron2.config import LazyConfig
|
|
|
|
from .visualize_humans import SemanticRenderer
|
|
from .. import REPO_PATH
|
|
|
|
# For Windows, remove PYOPENGL_PLATFORM to enable default rendering backend
|
|
sys_name = platform.system()
|
|
if sys_name == "Windows":
|
|
os.environ.pop("PYOPENGL_PLATFORM")
|
|
|
|
|
|
LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353)
|
|
|
|
|
|
def predict_smpl(batch, model, model_cfg, figure_scale=None):
|
|
all_verts = []
|
|
all_cam_t = []
|
|
with torch.no_grad():
|
|
out = model(batch)
|
|
|
|
pred_cam = out["pred_cam"]
|
|
pred_smpl_parameter = out["pred_smpl_params"]
|
|
if figure_scale is not None:
|
|
pred_smpl_parameter['betas'][0][1] = float(figure_scale)
|
|
smpl_output = model.smpl(
|
|
**{k: v.float() for k, v in pred_smpl_parameter.items()},
|
|
pose2rot=False,
|
|
)
|
|
pred_vertices = smpl_output.vertices
|
|
out["pred_vertices"] = pred_vertices.reshape(
|
|
batch["img"].shape[0], -1, 3
|
|
)
|
|
|
|
box_center = batch["box_center"].float()
|
|
box_size = batch["box_size"].float()
|
|
img_size = batch["img_size"].float()
|
|
|
|
scaled_focal_length = (
|
|
model_cfg.EXTRA.FOCAL_LENGTH
|
|
/ model_cfg.MODEL.IMAGE_SIZE
|
|
* img_size.max()
|
|
)
|
|
pred_cam_t_full = cam_crop_to_full(
|
|
pred_cam, box_center, box_size, img_size, scaled_focal_length
|
|
).detach().cpu().numpy()
|
|
# Render the result
|
|
batch_size = batch["img"].shape[0]
|
|
for n in range(batch_size):
|
|
# Add all verts and cams to list
|
|
verts = out["pred_vertices"][n].detach().cpu().numpy()
|
|
cam_t = pred_cam_t_full[n]
|
|
all_verts.append(verts)
|
|
all_cam_t.append(cam_t)
|
|
|
|
misc_args = dict(
|
|
mesh_base_color=LIGHT_BLUE,
|
|
scene_bg_color=(1, 1, 1),
|
|
focal_length=scaled_focal_length,
|
|
)
|
|
|
|
smpl_outs = {
|
|
k: v.detach().cpu().numpy() for k, v in pred_smpl_parameter.items()
|
|
}
|
|
|
|
results_dict_for_rendering = {
|
|
"verts": all_verts,
|
|
"cam_t": all_cam_t,
|
|
"render_res": img_size[n].cpu().numpy(),
|
|
"smpls": smpl_outs,
|
|
"scaled_focal_length": scaled_focal_length.cpu().numpy(),
|
|
}
|
|
return results_dict_for_rendering, misc_args
|
|
|
|
|
|
def load_image(model_cfg, image, detector):
|
|
# Detect humans in image
|
|
# image = image.permute(1,2,0)*255
|
|
det_out = detector(image)
|
|
# image = image.numpy()
|
|
det_instances = det_out["instances"]
|
|
valid_idx = (det_instances.pred_classes == 0) & (det_instances.scores > 0.5)
|
|
boxes = det_instances.pred_boxes.tensor[valid_idx].cpu().numpy()
|
|
# Run HMR2.0 on all detected humans
|
|
dataset = ViTDetDataset(model_cfg, image, boxes)
|
|
return torch.utils.data.DataLoader(
|
|
dataset, batch_size=8, shuffle=False, num_workers=0
|
|
)
|
|
|
|
def load_hmr_model(hmr2_checkpoint_path, smpl_dir_path):
|
|
model, model_cfg = load_hmr2(hmr2_checkpoint_path,smpl_dir_path)
|
|
return {
|
|
'model': model,
|
|
'model_cfg': model_cfg
|
|
}
|
|
|
|
|
|
def load_detectron(detectron_checkpoint_path):
|
|
cfg_path = os.path.join(REPO_PATH, "configs", "cascade_mask_rcnn_vitdet_h_75ep.py")
|
|
detectron2_cfg = LazyConfig.load(str(cfg_path))
|
|
detectron2_cfg.train.init_checkpoint = detectron_checkpoint_path
|
|
for i in range(3):
|
|
detectron2_cfg.model.roi_heads.box_predictors[i].test_score_thresh = 0.25
|
|
detector = DefaultPredictor_Lazy(detectron2_cfg)
|
|
return detector
|
|
|
|
|
|
def process_humans(images, hmr, detector, figure_scale, device):
|
|
images = images.permute(0, 3, 1, 2)
|
|
_, _, height, width = images.shape
|
|
model, model_cfg = hmr['model'], hmr['model_cfg']
|
|
|
|
|
|
model = model.to(device)
|
|
detector.model.to(device)
|
|
# This PyRender is only used for visualizing, we use Blender after to render different conditions.
|
|
renderer = SemanticRenderer(
|
|
model_cfg,
|
|
faces=model.smpl.faces,
|
|
lbs=model.smpl.lbs_weights,
|
|
viewport_size=(720, 720),
|
|
)
|
|
results = []
|
|
rendered_images = []
|
|
semantic_maps = []
|
|
masks = []
|
|
focal_length = None
|
|
mesh_base_color = None
|
|
scene_bg_color = None
|
|
pbar = comfy.utils.ProgressBar(len(images))
|
|
for image in tqdm(images, desc="Processing Human Images:"):
|
|
image = image.permute(1,2,0)*255
|
|
image = image.numpy()
|
|
renderer.renderer.delete()
|
|
renderer.renderer = pyrender.OffscreenRenderer(
|
|
viewport_width=width,
|
|
viewport_height=height,
|
|
point_size=1.0,
|
|
)
|
|
dataloader = load_image(model_cfg, image, detector)
|
|
|
|
|
|
for batch in dataloader:
|
|
batch = recursive_to(batch, device)
|
|
results_dict_for_rendering, misc_args = predict_smpl(batch, model, model_cfg, figure_scale)
|
|
mesh_base_color = misc_args['mesh_base_color']
|
|
scene_bg_color = misc_args['scene_bg_color']
|
|
focal_length = misc_args['focal_length'].to('cpu')
|
|
|
|
rendering_results = renderer.render_all_multiple(
|
|
results_dict_for_rendering["verts"],
|
|
cam_t=results_dict_for_rendering["cam_t"],
|
|
render_res=results_dict_for_rendering["render_res"], **misc_args
|
|
)
|
|
# Overlay image
|
|
valid_mask = rendering_results["Image"][:, :, -1][:, :, np.newaxis]
|
|
cam_view = (
|
|
valid_mask * rendering_results["Image"][:, :, [2, 1, 0]]
|
|
+ (1 - valid_mask) * image.astype(np.float32)[:, :, ::-1] / 255
|
|
)
|
|
rendered_images.append(torch.from_numpy(cam_view[:, :, ::-1].copy()))
|
|
semantic_maps.append(torch.from_numpy(rendering_results.get("SemanticMap")))
|
|
masks.append(torch.from_numpy(rendering_results.get("Mask")[:, :, 0]))
|
|
results.append(results_dict_for_rendering) # still in np
|
|
pbar.update(1)
|
|
|
|
smpl = {
|
|
'results': results,
|
|
'width': width,
|
|
'height': height,
|
|
'model_cfg': model_cfg,
|
|
'faces': model.smpl.faces,
|
|
'lbs': model.smpl.lbs_weights,
|
|
'viewport_size': (720, 720),
|
|
'misc_args': { 'mesh_base_color': mesh_base_color, 'scene_bg_color': scene_bg_color, 'focal_length': focal_length }
|
|
}
|
|
|
|
return {
|
|
'smpl': smpl,
|
|
'images': torch.stack(rendered_images),
|
|
'masks': torch.stack(masks).to(torch.float16),
|
|
'semantic_maps': torch.stack(semantic_maps),
|
|
}
|
|
|
|
|
|
def get_human(images, smpl, idx):
|
|
results = []
|
|
rendered_images = []
|
|
semantic_maps = []
|
|
masks = []
|
|
|
|
model_cfg, faces, lbs = smpl['model_cfg'], smpl['faces'], smpl['lbs']
|
|
viewport_size = smpl['viewport_size']
|
|
width, height = smpl['width'], smpl['height']
|
|
misc_args = smpl['misc_args']
|
|
|
|
renderer = SemanticRenderer(
|
|
model_cfg,
|
|
faces=faces,
|
|
lbs=lbs,
|
|
viewport_size=viewport_size,
|
|
)
|
|
|
|
renderer.renderer = pyrender.OffscreenRenderer(
|
|
viewport_width=width,
|
|
viewport_height=height,
|
|
point_size=1.0,
|
|
)
|
|
|
|
human_smpl = []
|
|
keys = ['verts']
|
|
smpl_keys = ['global_orient', 'body_pose', 'betas']
|
|
for results_dict_for_rendering in smpl['results']:
|
|
result = {
|
|
'scaled_focal_length': results_dict_for_rendering['scaled_focal_length'],
|
|
'smpls': {},
|
|
'cam_t': results_dict_for_rendering['cam_t'],
|
|
'render_res': results_dict_for_rendering['render_res']
|
|
}
|
|
good = True
|
|
for key in keys:
|
|
if idx < len(results_dict_for_rendering[key]):
|
|
result[key] = [results_dict_for_rendering[key][idx]]
|
|
else:
|
|
good = False
|
|
|
|
for key in smpl_keys:
|
|
if idx < len(results_dict_for_rendering['smpls'][key]):
|
|
result['smpls'][key] = [results_dict_for_rendering['smpls'][key][idx]]
|
|
else:
|
|
good = False
|
|
|
|
if good:
|
|
human_smpl.append(result)
|
|
else:
|
|
human_smpl.append(None)
|
|
|
|
for image_orig, results_dict_for_rendering in zip(images, human_smpl):
|
|
image = image_orig.permute(1,2,0)*255
|
|
image = image.numpy()
|
|
if results_dict_for_rendering is not None:
|
|
rendering_results = renderer.render_all_multiple(
|
|
results_dict_for_rendering["verts"],
|
|
cam_t=results_dict_for_rendering["cam_t"],
|
|
render_res=results_dict_for_rendering["render_res"], **misc_args
|
|
)
|
|
# Overlay image
|
|
valid_mask = rendering_results["Image"][:, :, -1][:, :, np.newaxis]
|
|
cam_view = (
|
|
valid_mask * rendering_results["Image"][:, :, [2, 1, 0]]
|
|
+ (1 - valid_mask) * image.astype(np.float32)[:, :, ::-1] / 255
|
|
)
|
|
rendered_images.append(torch.from_numpy(cam_view[:, :, ::-1].copy()))
|
|
semantic_maps.append(torch.from_numpy(rendering_results.get("SemanticMap")))
|
|
masks.append(torch.from_numpy(rendering_results.get("Mask")[:, :, 0]))
|
|
else:
|
|
rendered_images.append(image_orig)
|
|
semantic_maps.append(torch.ones_like(image_orig))
|
|
masks.append(torch.zeros_like(image_orig))
|
|
results.append(results_dict_for_rendering) # still in np
|
|
|
|
smpl = {
|
|
'results': results,
|
|
'width': width,
|
|
'height': height,
|
|
'model_cfg': model_cfg,
|
|
'faces': faces,
|
|
'lbs': lbs,
|
|
'viewport_size': (720, 720),
|
|
'misc_args': misc_args
|
|
}
|
|
|
|
return {
|
|
'smpl': smpl,
|
|
'images': torch.stack(rendered_images),
|
|
'masks': torch.stack(masks).to(torch.float16),
|
|
'semantic_maps': torch.stack(semantic_maps),
|
|
}
|