【玩轉Colab】GitHub標星3.6k | 給AI一張高清照片，分分鐘還你細節滿滿的3D人體模型！

圖文介紹部分轉自公眾號：計算機視覺Daily

傳送門

GitHub地址：
https://github.com/facebookresearch/pifuhd

Demo地址：
https://colab.research.google.com/drive/11z58bl3meSzo6kFqkahMa35G5jmh2Wgt?sp=sharing#scrollTo=afwL_-ROCmDf

圖文介紹

手動對人體進行3D建模並非易事。

但現在，只給AI一張高清照片，它還真就能分分鐘搞定這件事。

甚至還挺高清，衣服褶皺、面部表情，細節一點不少。

這項新研究來自南加州大學和Facebook，中選CVPR 2020。

並且已經在GitHub上開源，標星3.6k，還在一天內就漲了207顆星，登上GitHub熱榜。

一起來看看，這究竟是如何實現的。

多級畫素對齊隱式函式

這隻AI名叫PIFuHD，其基礎框架是ICCV 2019上已經登場的畫素對齊隱式函式PIFu。不過，PIFu以解析度為512×512的影像作為輸入，輸出的3D模型解析度不高

為了得到高解析度的輸出，在這項研究中，研究人員在PIFu的基礎之上，額外疊加了一個畫素對齊的預測模組。

如圖所示，頂部粗層次畫素對齊預測器捕捉全域性的3D結構。高解析度的細節則由下面的Fine模組新增。

具體而言，fine模組將1024×1024的影像作為輸入，並將其編碼成高解析度的影像特徵（512×512）。

此後，高解析度特徵嵌入和第一個模組中得到的3D嵌入被結合起來，用以預測佔位概率場。

為了進一步提高重建的質量和保真度，該方法還會在影像空間中預測正反兩面的法線圖，並將其作為額外的輸入反饋給網路。

Demo可玩

論文程式碼已經開源，並且，研究團隊還在Colab上提供了線上試玩。

輸入一張你自己的照片，幾分鐘之內就能收穫一個數字3D的你。

真·3D建模師福音。

結合可以讓3D模型動起來的Mixamo食用，網友們都玩嗨了。

執行程式碼

Colab連結：
https://colab.research.google.com/drive/11z58bl3meSzo6kFqkahMa35G5jmh2Wgt?sp=sharing#scrollTo=afwL_-ROCmDf

進入Colab，將其拷貝至自己的雲端硬碟裡。

進入自己的雲端硬碟，找到剛才拷貝過來的副本，點選進去。

因為Colab的Pytorch版本不對，需要將pytorch版本更新為1.6.0

先執行程式碼：

!pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

依次執行程式碼框：

Clone PIFuHD repository

!git clone https://github.com/facebookresearch/pifuhd

Configure input data

cd /content/pifuhd/sample_images

from google.colab import files

filename = list(files.upload().keys())[0]

這裡需要選擇本地圖片，限制512*512畫素的圖片才可以，如果尺寸不對，就先裁剪或resize一下。

import os

try:
  image_path = '/content/pifuhd/sample_images/%s' % filename
except:
  image_path = '/content/pifuhd/sample_images/test.png' # example image
image_dir = os.path.dirname(image_path)
file_name = os.path.splitext(os.path.basename(image_path))[0]

# output pathes
obj_path = '/content/pifuhd/results/pifuhd_final/recon/result_%s_256.obj' % file_name
out_img_path = '/content/pifuhd/results/pifuhd_final/recon/result_%s_256.png' % file_name
video_path = '/content/pifuhd/results/pifuhd_final/recon/result_%s_256.mp4' % file_name
video_display_path = '/content/pifuhd/results/pifuhd_final/result_%s_256_display.mp4' % file_name

cd /content

Preprocess (for cropping image)

!git clone https://github.com/Daniil-Osokin/lightweight-human-pose-estimation.pytorch.git

cd /content/lightweight-human-pose-estimation.pytorch/

!wget https://download.01.org/opencv/openvino_training_extensions/models/human_pose_estimation/checkpoint_iter_370000.pth

import torch
import cv2
import numpy as np
from models.with_mobilenet import PoseEstimationWithMobileNet
from modules.keypoints import extract_keypoints, group_keypoints
from modules.load_state import load_state
from modules.pose import Pose, track_poses
import demo

def get_rect(net, images, height_size):
    net = net.eval()

    stride = 8
    upsample_ratio = 4
    num_keypoints = Pose.num_kpts
    previous_poses = []
    delay = 33
    for image in images:
        rect_path = image.replace('.%s' % (image.split('.')[-1]), '_rect.txt')
        img = cv2.imread(image, cv2.IMREAD_COLOR)
        orig_img = img.copy()
        orig_img = img.copy()
        heatmaps, pafs, scale, pad = demo.infer_fast(net, img, height_size, stride, upsample_ratio, cpu=False)

        total_keypoints_num = 0
        all_keypoints_by_type = []
        for kpt_idx in range(num_keypoints):  # 19th for bg
            total_keypoints_num += extract_keypoints(heatmaps[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num)

        pose_entries, all_keypoints = group_keypoints(all_keypoints_by_type, pafs, demo=True)
        for kpt_id in range(all_keypoints.shape[0]):
            all_keypoints[kpt_id, 0] = (all_keypoints[kpt_id, 0] * stride / upsample_ratio - pad[1]) / scale
            all_keypoints[kpt_id, 1] = (all_keypoints[kpt_id, 1] * stride / upsample_ratio - pad[0]) / scale
        current_poses = []

        rects = []
        for n in range(len(pose_entries)):
            if len(pose_entries[n]) == 0:
                continue
            pose_keypoints = np.ones((num_keypoints, 2), dtype=np.int32) * -1
            valid_keypoints = []
            for kpt_id in range(num_keypoints):
                if pose_entries[n][kpt_id] != -1.0:  # keypoint was found
                    pose_keypoints[kpt_id, 0] = int(all_keypoints[int(pose_entries[n][kpt_id]), 0])
                    pose_keypoints[kpt_id, 1] = int(all_keypoints[int(pose_entries[n][kpt_id]), 1])
                    valid_keypoints.append([pose_keypoints[kpt_id, 0], pose_keypoints[kpt_id, 1]])
            valid_keypoints = np.array(valid_keypoints)
            
            if pose_entries[n][10] != -1.0 or pose_entries[n][13] != -1.0:
              pmin = valid_keypoints.min(0)
              pmax = valid_keypoints.max(0)

              center = (0.5 * (pmax[:2] + pmin[:2])).astype(np.int)
              radius = int(0.65 * max(pmax[0]-pmin[0], pmax[1]-pmin[1]))
            elif pose_entries[n][10] == -1.0 and pose_entries[n][13] == -1.0 and pose_entries[n][8] != -1.0 and pose_entries[n][11] != -1.0:
              # if leg is missing, use pelvis to get cropping
              center = (0.5 * (pose_keypoints[8] + pose_keypoints[11])).astype(np.int)
              radius = int(1.45*np.sqrt(((center[None,:] - valid_keypoints)**2).sum(1)).max(0))
              center[1] += int(0.05*radius)
            else:
              center = np.array([img.shape[1]//2,img.shape[0]//2])
              radius = max(img.shape[1]//2,img.shape[0]//2)

            x1 = center[0] - radius
            y1 = center[1] - radius

            rects.append([x1, y1, 2*radius, 2*radius])

        np.savetxt(rect_path, np.array(rects), fmt='%d')

net = PoseEstimationWithMobileNet()
checkpoint = torch.load('checkpoint_iter_370000.pth', map_location='cpu')
load_state(net, checkpoint)

get_rect(net.cuda(), [image_path], 512)

Preprocess (for cropping image)

cd /content/pifuhd/

!sh ./scripts/download_trained_model.sh

Run PIFuHD!

# Warning: all images with the corresponding rectangle files under -i will be processed. 
!python -m apps.simple_test -r 256 --use_rect -i $image_dir

# seems that 256 is the maximum resolution that can fit into Google Colab. 
# If you want to reconstruct a higher-resolution mesh, please try with your own machine. 

Render the result

# This command takes a few minutes
!pip install pytorch3d

!pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

from lib.colab_util import generate_video_from_obj, set_renderer, video

renderer = set_renderer()
generate_video_from_obj(obj_path, out_img_path, video_path, renderer)

# we cannot play a mp4 video generated by cv2
!ffmpeg -i $video_path -vcodec libx264 $video_display_path -y -loglevel quiet
video(video_display_path)

我的測試效果