計算機視覺人體例項分割之pose2seg

程式碼連線：

文章連線：

cvpr2019新出的人體例項分割模型，效果還不錯，就是得需要提取人體關鍵點，比起maskrcnn要麻煩一點，但效果要比maskrcnn好，尤其是人員者當時。開源的程式碼中沒有預測的程式碼，對其test.py進行了修改，分割結果就可以視覺化了。

在k80上測試，單張圖片平均耗時0.297秒，高階顯示卡會更快些。

import argparse

import numpy as np

from tqdm import tqdm

from modeling.build_model import Pose2Seg

from datasets.CocoDatasetInfo import CocoDatasetInfo, annToMask

from pycocotools import mask as maskUtils

import cv2

import time

def apply_mask(image, mask, color, alpha=0.5):

 for c in range(3):

 image[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - alpha) + alpha * color[c] * 255, image[:, :, c])

 return image

def predict(model, dataset='cocoVal', logger=print):

 if dataset == 'OCHumanVal':

 ImageRoot = './data/OCHuman/images'

 AnnoFile = './data/OCHuman/annotations/ochuman_coco_format_val_range_0.00_1.00.json'

 elif dataset == 'OCHumanTest':

 ImageRoot = './data/OCHuman/images'

 AnnoFile = './data/OCHuman/annotations/ochuman_coco_format_test_range_0.00_1.00.json'

 elif dataset == 'cocoVal':

 ImageRoot = './data/val2017'

 AnnoFile = './data/annotations/person_keypoints_val2017_pose2seg.json'

 datainfos = CocoDatasetInfo(ImageRoot, AnnoFile, loadimg=True)

 model.eval()

 results_segm = []

 imgIds = []

 total_time = 0.

 for i in tqdm(range(len(datainfos))):

 # for i in tqdm(range(1)):

 rawdata = datainfos[i]

 img = rawdata['data']

 image_id = rawdata['id']

 height, width = img.shape[0:2]

 gt_kpts = np.float32(rawdata['gt_keypoints']).transpose(0, 2, 1) # (N, 17, 3)

 # print(gt_kpts)

 # gt_segms = rawdata['segms']

 # gt_masks = np.array([annToMask(segm, height, width) for segm in gt_segms])

 start_time = time.time()

 output = model([img], [gt_kpts])

 end_time = time.time()

 total_time = total_time + end_time - start_time

 for mask in output[0]:

 contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

 cv2.drawContours(img, contours, -1,

 (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255)), 3)

 color = [np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255)]

 img = apply_mask(img, mask, color, 0.9)

 maskencode = maskUtils.encode(np.asfortranarray(mask))

 maskencode['counts'] = maskencode['counts'].decode('ascii')

 results_segm.append({

 "image_id": image_id,

 "category_id": 1,

 "score": 1.0,

 "segmentation": maskencode

 })

 cv2.imwrite('./figures/mask' + str(i) + '.jpg', img)

 imgIds.append(image_id)

 print('============time: ', total_time * 1.0 / len(datainfos))

if __name__ == '__main__':

 parser = argparse.ArgumentParser(description="Pose2Seg Testing")

 parser.add_argument(

 "--weights",

 help="path to .pkl model weight",

 type=str,

 )

 parser.add_argument(

 "--coco",

 default=True,

 help="Do test on COCOPersons val set",

 action="store_true",

 )

 parser.add_argument(

 "--OCHuman",

 help="Do test on OCHuman val&test set",

 action="store_true",

 )

 args = parser.parse_args()

 print('===========> loading model <===========')

 model = Pose2Seg().cuda()

 model.init(args.weights)

 print('===========> testing <===========')

 if args.coco:

 predict(model, dataset='cocoVal')

 if args.OCHuman:

 predict(model, dataset='OCHumanVal')

 predict(model, dataset='OCHumanTest')