Spaces:
Running
Running
File size: 5,381 Bytes
48cafca 935680d 48cafca 935680d 48cafca 5fdb397 48cafca 5fdb397 48cafca 8f8e463 48cafca |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
from pathlib import Path
import torch
import os
import cv2
import numpy as np
import tempfile
from tqdm import tqdm
import torch.utils
import trimesh
import torch.utils.data
import gradio as gr
from typing import Union, List, Tuple, Dict
from amr.models import AMR
from amr.configs import get_config
from amr.utils import recursive_to
from amr.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD
from amr.utils.renderer import Renderer, cam_crop_to_full
from huggingface_hub import snapshot_download
LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353)
# Load model config
path_model_cfg = 'config/config.yaml'
model_cfg = get_config(path_model_cfg)
# Load model
repo_id = "luoxue-star/AniMer"
local_dir = snapshot_download(repo_id=repo_id)
# local_dir = "./checkpoints"
PATH_CHECKPOINT = os.path.join(local_dir, "checkpoint.ckpt")
model = AMR.load_from_checkpoint(checkpoint_path=PATH_CHECKPOINT, map_location="cpu",
cfg=model_cfg, strict=False, weights_only=True)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = model.to(device)
model.eval()
# Setup the renderer
renderer = Renderer(model_cfg, faces=model.smal.faces)
# Make output directory if it does not exist
OUTPUT_FOLDER = "demo_out"
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
def predict(im):
return im["composite"]
def inference(img: Dict)-> Tuple[Union[np.ndarray|None], List[str]]:
img = np.array(img["composite"])[:, :, :-1]
boxes = np.array([[0, 0, img.shape[1], img.shape[0]]]) # x1, y1, x2, y2
# Run AniMer on the crop image
dataset = ViTDetDataset(model_cfg, img, boxes)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8)
all_verts = []
all_cam_t = []
temp_name = next(tempfile._get_candidate_names())
for batch in tqdm(dataloader):
batch = recursive_to(batch, device)
with torch.no_grad():
out = model(batch)
pred_cam = out['pred_cam']
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):
person_id = int(batch['personid'][n])
input_patch = (batch['img'][n].cpu() * 255 * (DEFAULT_STD[:, None, None]) + (
DEFAULT_MEAN[:, None, None])) / 255.
input_patch = input_patch.permute(1, 2, 0).numpy()
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)
# Render mesh onto the original image
if len(all_verts):
misc_args = dict(
mesh_base_color=LIGHT_BLUE,
scene_bg_color=(1, 1, 1),
focal_length=scaled_focal_length,
)
cam_view = renderer.render_rgba_multiple(all_verts, cam_t=all_cam_t, render_res=img_size[n], **misc_args)
# Overlay image
input_img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR).astype(np.float32)[:, :, ::-1] / 255.0
input_img = np.concatenate([input_img, np.ones_like(input_img[:, :, :1])], axis=2) # Add alpha channel
input_img_overlay = input_img[:, :, :3] * (1 - cam_view[:, :, 3:]) + cam_view[:, :, :3] * cam_view[:, :, 3:]
output_img = (255 * input_img_overlay[:, :, ::-1]).astype(np.uint8)[:, :, [2, 1, 0]]
# Return mesh path
trimeshes = [renderer.vertices_to_trimesh(vvv, ttt.copy(), LIGHT_BLUE) for vvv,ttt in zip(all_verts, all_cam_t)]
# Join meshes
mesh = trimesh.util.concatenate(trimeshes)
# Save mesh to file
mesh_name = os.path.join(OUTPUT_FOLDER, next(tempfile._get_candidate_names()) + '.obj')
trimesh.exchange.export.export_mesh(mesh, mesh_name)
return (output_img, mesh_name)
else:
return (None, [])
demo = gr.Interface(
fn=inference,
analytics_enabled=False,
inputs=gr.ImageEditor(label="Input image", sources=["upload", "clipboard"], type='pil',
brush=False, eraser=False, layers=False, transforms="crop",
interactive=True),
outputs=[
gr.Image(label="Overlap image"),
gr.Model3D(display_mode="wireframe", label="3D Mesh"),
],
title="AniMer: Animal Pose and Shape Estimation Using Family Aware Transformer",
description="""
Project page: https://luoxue-star.github.io/AniMer_project_page/
## Steps for Use
1. **Input**: Select an example image or upload your own image.
2. **Crop**: Crop the animal in the image (Otherwise, the result may be poor.)
3. **Output**:
- Overlapping Image
- 3D Mesh
""",
examples=[
'example_data/000000015956_horse.png',
'example_data/n02101388_1188.png',
'example_data/n02412080_12159.png',
'example_data/000000101684_zebra.png',
],
)
demo.launch()
|