get colored uv map for SMPL/SMPLX model

smplx_uv.py

# -*- coding: utf-8 -*-
"""
Created on Sat Oct 10 15:29:24 2020
@author: qzane
The SMPLX models can be find at: https://smpl-x.is.tue.mpg.de/
"""

import numpy as np
import scipy
from scipy.ndimage.measurements import label
import cv2


input_obj = './smplx_uv/smplx_uv.obj'
output_texture = './smplx_uv/smplx_uv_color_texture.png'
texture_size = 1024


def cv2_triangle(img, p123):
    p1, p2, p3 = (tuple(i) for i in p123)
    cv2.line(img, p1, p2, (255, 0, 0), 1) 
    cv2.line(img, p2, p3, (255, 0, 0), 1) 
    cv2.line(img, p1, p3, (255, 0, 0), 1)
    return img

def obj_vt(fname=input_obj): #vertices
    res = []
    with open(fname) as f:
        for line in f:
            if line.startswith('vt '):
                tmp = line.split(' ')
                v = [float(i) for i in tmp[1:3]]
                res.append(v)
    return np.array(res, dtype=np.float)

def obj_ft(fname=input_obj): #faces: vt,2,3
    res = []
    with open(fname) as f:
        for line in f:
            if line.startswith('f '):
                tmp = line.split(' ')
                if '/' in tmp[1]:
                    v = [int(i.split('/')[1]) for i in tmp[1:4]]
                else:
                    raise(Exception("not a vallid obj file"))
                res.append(v)
    return np.array(res, dtype=np.int) - 1 # obj index from 1


vt = obj_vt(input_obj).reshape(-1, 2)
vt[:, 1] = 1 - vt[:, 1] # texture origin is on the bottom left
ft = obj_ft(input_obj)

vt = (vt * texture_size).reshape(-1, 2).astype(np.int32) # cv2.drawContours requires int32
texture0 = np.zeros((texture_size, texture_size, 3), np.uint8)

# draw without colors
for v123 in ft:
    cv2.drawContours(texture0, [vt[v123]], 0, (255,255,255), -1)

# get labels    
labels, num_label = label(texture0[:,:,0])

colors = np.arange(0,13).astype(np.uint8)*19
colors = cv2.applyColorMap(colors, cv2.COLORMAP_HSV)
colors = colors.reshape(-1, 3).tolist()

# draw colors
for i in range(1, num_label+1):
    texture0[labels==i] = colors[i%13]   
    
# draw triangles    
for t1,t2,t3 in ft:
    p1 = (vt[t1]).tolist()
    p2 = (vt[t2]).tolist()
    p3 = (vt[t3]).tolist()
    cv2_triangle(texture0, (p1,p2,p3))
    
# draw ids  
for i in range(1, num_label+1):
    q = np.where(labels==i)
    q = np.array(q)
    x,y = q.mean(1).astype(np.int32)
    cv2.putText(texture0, str(i), (y,x), cv2.FONT_HERSHEY_COMPLEX_SMALL, 2, (0,0,0), 3) 
    cv2.putText(texture0, str(i), (y,x), cv2.FONT_HERSHEY_COMPLEX_SMALL, 2, (255,255,255), 1) 

cv2.imwrite(output_texture, texture0)