First, you need to use dump
command to generate pkl
file after that we can generate the needed IUV image.
Inside Densepose use apply_net.py
to generate the pkl
file. It has a lot of other options as well. You can check them out here.
python3 apply_net.py dump configs/densepose_rcnn_R_101_FPN_s1x.yaml /models/R_101_FPN_s1x.pkl /Images/frame.jpg --output output.pkl -v
After having the pkl
file we will need to extract information from it and create your IUV needed image.
img = Image.open('/Images/frame.jpg')
img_w ,img_h = img.size
with open('output.pkl','rb') as f:
data=pickle.load(f)
i = data[0]['pred_densepose'][0].labels.cpu().numpy()
uv = data[0]['pred_densepose'][0].uv.cpu().numpy()
iuv = np.stack((uv[1,:,:], uv[0,:,:], i * 0,))
iuv = np.transpose(iuv, (1,2,0))
iuv_img = Image.fromarray(np.uint8(iuv*255),"RGB")
iuv_img.show() #It shows only the croped person
box = data[0]["pred_boxes_XYXY"][0]
box[2]=box[2]-box[0]
box[3]=box[3]-box[1]
x,y,w,h=[int(v) for v in box]
bg=np.zeros((img_h,img_w,3))
bg[y:y+h,x:x+w,:]=iuv
bg_img = Image.fromarray(np.uint8(bg*255),"RGB")
bg_img.save('output.png')