release load and visualize MPII dataset

zsdonghao · zsdonghao · commit d10d1060287c · 2018-04-10T19:29:30.000+01:00
diff --git a/docs/modules/files.rst b/docs/modules/files.rst
@@ -20,6 +20,7 @@ API - Files
    load_cyclegan_dataset
    load_celebA_dataset
    load_voc_dataset
+   load_mpii_pose_dataset
    download_file_from_google_drive
 
    save_npz
@@ -108,6 +109,10 @@ VOC 2007/2012
 ^^^^^^^^^^^^^^^^
 .. autofunction:: load_voc_dataset
 
+MPII 
+^^^^^^^^^^^^^^^^
+.. autofunction:: load_mpii_pose_dataset
+
 Google Drive
 ^^^^^^^^^^^^^^^^
 .. autofunction:: download_file_from_google_drive
diff --git a/docs/modules/visualize.rst b/docs/modules/visualize.rst
@@ -13,6 +13,7 @@ to visualize the model, activations etc. Here we provide more functions for data
    save_image
    save_images
    draw_boxes_and_labels_to_image
+   draw_mpii_people_to_image
    draw_weights
    CNN2d
    frame
@@ -44,6 +45,9 @@ Save image for object detection
 .. autofunction:: draw_boxes_and_labels_to_image
 
 
+Save image for pose estimation (MPII)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: draw_mpii_people_to_image
 
 Visualize model parameters
 ------------------------------
diff --git a/tensorlayer/files.py b/tensorlayer/files.py
@@ -65,6 +65,7 @@
     'download_file_from_google_drive',
     'load_celebA_dataset',
     'load_voc_dataset',
+    'load_mpii_pose_dataset',
     'save_npz',
     'load_npz',
     'assign_params',
@@ -1317,6 +1318,232 @@ def convert_annotation(file_name):
         n_objs_list, objs_info_list, objs_info_dicts
 
 
+def load_mpii_pose_dataset(path='data', is_16_pos_only=False):
+    """Load MPII Human Pose Dataset.
+
+    Parameters
+    -----------
+    path : str
+        The path that the data is downloaded to.
+    is_16_pos_only : boolean
+        If True, only return the peoples contain 16 pose keypoints. (Usually be used for single person pose estimation)
+
+    Returns
+    ----------
+    img_train_list : list of str
+        The image directories of training data.
+    ann_train_list : list of dict
+        The annotations of training data.
+    img_test_list : list of str
+        The image directories of testing data.
+    ann_test_list : list of dict
+        The annotations of testing data.
+
+    Examples
+    --------
+    >>> import pprint
+    >>> import tensorlayer as tl
+    >>> img_train_list, ann_train_list, img_test_list, ann_test_list = tl.files.load_mpii_pose_dataset()
+    >>> image = tl.vis.read_image(img_train_list[0])
+    >>> tl.vis.draw_mpii_people_to_image(image, ann_train_list[0], 'image.png')
+    >>> pprint.pprint(ann_train_list[0])
+
+    References
+    -----------
+    - `MPII Human Pose Dataset. CVPR 14 <http://human-pose.mpi-inf.mpg.de>`__
+    - `MPII Human Pose Models. CVPR 16 <http://pose.mpi-inf.mpg.de>`__
+    - `MPII Human Shape, Poselet Conditioned Pictorial Structures and etc <http://pose.mpi-inf.mpg.de/#related>`__
+    - `MPII Keyponts and ID <http://human-pose.mpi-inf.mpg.de/#download>`__
+    """
+    path = os.path.join(path, 'mpii_human_pose')
+    logging.info("Load or Download MPII Human Pose > {}".format(path))
+
+    # annotation
+    url = "http://datasets.d2.mpi-inf.mpg.de/andriluka14cvpr/"
+    tar_filename = "mpii_human_pose_v1_u12_2.zip"
+    extracted_filename = "mpii_human_pose_v1_u12_2"
+    if folder_exists(os.path.join(path, extracted_filename)) is False:
+        logging.info("[MPII] (annotation) {} is nonexistent in {}".format(extracted_filename, path))
+        maybe_download_and_extract(tar_filename, path, url, extract=True)
+        del_file(os.path.join(path, tar_filename))
+
+    # images
+    url = "http://datasets.d2.mpi-inf.mpg.de/andriluka14cvpr/"
+    tar_filename = "mpii_human_pose_v1.tar.gz"
+    extracted_filename2 = "images"
+    if folder_exists(os.path.join(path, extracted_filename2)) is False:
+        logging.info("[MPII] (images) {} is nonexistent in {}".format(extracted_filename, path))
+        maybe_download_and_extract(tar_filename, path, url, extract=True)
+        del_file(os.path.join(path, tar_filename))
+
+    # parse annotation, format see http://human-pose.mpi-inf.mpg.de/#download
+    import scipy.io as sio
+    logging.info("reading annotations from mat file ...")
+    # mat = sio.loadmat(os.path.join(path, extracted_filename, "mpii_human_pose_v1_u12_1.mat"))
+
+    # def fix_wrong_joints(joint):    # https://github.com/mitmul/deeppose/blob/master/datasets/mpii_dataset.py
+    #     if '12' in joint and '13' in joint and '2' in joint and '3' in joint:
+    #         if ((joint['12'][0] < joint['13'][0]) and
+    #                 (joint['3'][0] < joint['2'][0])):
+    #             joint['2'], joint['3'] = joint['3'], joint['2']
+    #         if ((joint['12'][0] > joint['13'][0]) and
+    #                 (joint['3'][0] > joint['2'][0])):
+    #             joint['2'], joint['3'] = joint['3'], joint['2']
+    #     return joint
+
+    ann_train_list = []
+    ann_test_list = []
+    img_train_list = []
+    img_test_list = []
+
+    def save_joints():
+        # joint_data_fn = os.path.join(path, 'data.json')
+        # fp = open(joint_data_fn, 'w')
+        mat = sio.loadmat(os.path.join(path, extracted_filename, "mpii_human_pose_v1_u12_1.mat"))
+
+        for i, (anno, train_flag) in enumerate(  # all images
+                zip(mat['RELEASE']['annolist'][0, 0][0], mat['RELEASE']['img_train'][0, 0][0])):
+
+            img_fn = anno['image']['name'][0, 0][0]
+            train_flag = int(train_flag)
+
+            # print(i, img_fn, train_flag) # DEBUG print all images
+
+            if train_flag:
+                img_train_list.append(img_fn)
+                ann_train_list.append([])
+            else:
+                img_test_list.append(img_fn)
+                ann_test_list.append([])
+
+            head_rect = []
+            if 'x1' in str(anno['annorect'].dtype):
+                head_rect = zip([x1[0, 0] for x1 in anno['annorect']['x1'][0]], [y1[0, 0] for y1 in anno['annorect']['y1'][0]],
+                                [x2[0, 0] for x2 in anno['annorect']['x2'][0]], [y2[0, 0] for y2 in anno['annorect']['y2'][0]])
+            else:
+                head_rect = []  # TODO
+
+            if 'annopoints' in str(anno['annorect'].dtype):
+                annopoints = anno['annorect']['annopoints'][0]
+                head_x1s = anno['annorect']['x1'][0]
+                head_y1s = anno['annorect']['y1'][0]
+                head_x2s = anno['annorect']['x2'][0]
+                head_y2s = anno['annorect']['y2'][0]
+                for annopoint, head_x1, head_y1, head_x2, head_y2 in zip(annopoints, head_x1s, head_y1s, head_x2s, head_y2s):
+                    if annopoint != []:
+                        head_rect = [float(head_x1[0, 0]), float(head_y1[0, 0]), float(head_x2[0, 0]), float(head_y2[0, 0])]
+
+                        # joint coordinates
+                        annopoint = annopoint['point'][0, 0]
+                        j_id = [str(j_i[0, 0]) for j_i in annopoint['id'][0]]
+                        x = [x[0, 0] for x in annopoint['x'][0]]
+                        y = [y[0, 0] for y in annopoint['y'][0]]
+                        joint_pos = {}
+                        for _j_id, (_x, _y) in zip(j_id, zip(x, y)):
+                            joint_pos[int(_j_id)] = [float(_x), float(_y)]
+                        # joint_pos = fix_wrong_joints(joint_pos)
+
+                        # visiblity list
+                        if 'is_visible' in str(annopoint.dtype):
+                            vis = [v[0] if v else [0] for v in annopoint['is_visible'][0]]
+                            vis = dict([(k, int(v[0])) if len(v) > 0 else v for k, v in zip(j_id, vis)])
+                        else:
+                            vis = None
+
+                        # if len(joint_pos) == 16:
+                        if ((is_16_pos_only == True) and (len(joint_pos) == 16)) or (is_16_pos_only == False):
+                            # only use image with 16 key points / or use all
+                            data = {'filename': img_fn, 'train': train_flag, 'head_rect': head_rect, 'is_visible': vis, 'joint_pos': joint_pos}
+                            # print(json.dumps(data), file=fp)  # py3
+                            if train_flag:
+                                ann_train_list[-1].append(data)
+                            else:
+                                ann_test_list[-1].append(data)
+
+    # def write_line(datum, fp):
+    #     joints = sorted([[int(k), v] for k, v in datum['joint_pos'].items()])
+    #     joints = np.array([j for i, j in joints]).flatten()
+    #
+    #     out = [datum['filename']]
+    #     out.extend(joints)
+    #     out = [str(o) for o in out]
+    #     out = ','.join(out)
+    #
+    #     print(out, file=fp)
+
+    # def split_train_test():
+    #     # fp_test = open('data/mpii/test_joints.csv', 'w')
+    #     fp_test = open(os.path.join(path, 'test_joints.csv'), 'w')
+    #     # fp_train = open('data/mpii/train_joints.csv', 'w')
+    #     fp_train = open(os.path.join(path, 'train_joints.csv'), 'w')
+    #     # all_data = open('data/mpii/data.json').readlines()
+    #     all_data = open(os.path.join(path, 'data.json')).readlines()
+    #     N = len(all_data)
+    #     N_test = int(N * 0.1)
+    #     N_train = N - N_test
+    #
+    #     print('N:{}'.format(N))
+    #     print('N_train:{}'.format(N_train))
+    #     print('N_test:{}'.format(N_test))
+    #
+    #     np.random.seed(1701)
+    #     perm = np.random.permutation(N)
+    #     test_indices = perm[:N_test]
+    #     train_indices = perm[N_test:]
+    #
+    #     print('train_indices:{}'.format(len(train_indices)))
+    #     print('test_indices:{}'.format(len(test_indices)))
+    #
+    #     for i in train_indices:
+    #         datum = json.loads(all_data[i].strip())
+    #         write_line(datum, fp_train)
+    #
+    #     for i in test_indices:
+    #         datum = json.loads(all_data[i].strip())
+    #         write_line(datum, fp_test)
+
+    save_joints()
+    # split_train_test()  #
+
+    ## read images dir
+    logging.info("reading images list ...")
+    img_dir = os.path.join(path, extracted_filename2)
+    _img_list = load_file_list(path=os.path.join(path, extracted_filename2), regx='\\.jpg', printable=False)
+    # ann_list = json.load(open(os.path.join(path, 'data.json')))
+    for i, im in enumerate(img_train_list):
+        if im not in _img_list:
+            print('missing training image {} in {} (remove from img(ann)_train_list)'.format(im, img_dir))
+            # img_train_list.remove(im)
+            del img_train_list[i]
+            del ann_train_list[i]
+    for i, im in enumerate(img_test_list):
+        if im not in _img_list:
+            print('missing testing image {} in {} (remove from img(ann)_test_list)'.format(im, img_dir))
+            # img_test_list.remove(im)
+            del img_train_list[i]
+            del ann_train_list[i]
+
+    ## check annotation and images
+    n_train_images = len(img_train_list)
+    n_test_images = len(img_test_list)
+    n_images = n_train_images + n_test_images
+    logging.info("n_images: {} n_train_images: {} n_test_images: {}".format(n_images, n_train_images, n_test_images))
+    n_train_ann = len(ann_train_list)
+    n_test_ann = len(ann_test_list)
+    n_ann = n_train_ann + n_test_ann
+    logging.info("n_ann: {} n_train_ann: {} n_test_ann: {}".format(n_ann, n_train_ann, n_test_ann))
+    n_train_people = len(sum(ann_train_list, []))
+    n_test_people = len(sum(ann_test_list, []))
+    n_people = n_train_people + n_test_people
+    logging.info("n_people: {} n_train_people: {} n_test_people: {}".format(n_people, n_train_people, n_test_people))
+    # add path to all image file name
+    for i in range(len(img_train_list)):
+        img_train_list[i] = os.path.join(img_dir, img_train_list[i])
+    for i in range(len(img_test_list)):
+        img_test_list[i] = os.path.join(img_dir, img_test_list[i])
+    return img_train_list, ann_train_list, img_test_list, ann_test_list
+
+
 def save_npz(save_list=None, name='model.npz', sess=None):
     """Input parameters and the file name, save parameters into .npz file. Use tl.utils.load_npz() to restore.
 
diff --git a/tensorlayer/visualize.py b/tensorlayer/visualize.py
@@ -18,6 +18,7 @@
     'save_image',
     'save_images',
     'draw_boxes_and_labels_to_image',
+    'draw_mpii_people_to_image',
     'frame',
     'CNN2d',
     'images2d',
@@ -221,6 +222,112 @@ def draw_boxes_and_labels_to_image(image, classes, coords, scores, classes_list,
     return image
 
 
+def draw_mpii_people_to_image(image, peoples, save_name='image.png'):
+    """Draw people(s) into image using MPII dataset format as input, return or save the result image.
+
+    This is an experimental API, can be changed in the future.
+
+    Parameters
+    -----------
+    image : numpy.array
+        The RGB image [height, width, channel].
+    people :
+    save_name : None or str
+        The name of image file (i.e. image.png), if None, not to save image.
+
+    Returns
+    --------
+    numpy.array
+        The saved image.
+
+    Examples
+    --------
+    >>> import pprint
+    >>> import tensorlayer as tl
+    >>> img_train_list, ann_train_list, img_test_list, ann_test_list = tl.files.load_mpii_pose_dataset()
+    >>> image = tl.vis.read_image(img_train_list[0])
+    >>> tl.vis.draw_mpii_people_to_image(image, ann_train_list[0], 'image.png')
+    >>> pprint.pprint(ann_train_list[0])
+
+    References
+    -----------
+    - `MPII Keyponts and ID <http://human-pose.mpi-inf.mpg.de/#download>`__
+    """
+    import cv2
+    # import skimage
+    # don't change the original image, and avoid error https://stackoverflow.com/questions/30249053/python-opencv-drawing-errors-after-manipulating-array-with-numpy
+    image = image.copy()
+    radius = int(image.shape[1] / 500) + 1
+
+    imh, imw = image.shape[0:2]
+    thick = int((imh + imw) // 430)
+
+    if image.max() < 1:
+        image = image * 255
+
+    for people in peoples:
+        joint_pos = people['joint_pos']
+        ## draw circles
+        for pos in joint_pos.items():
+            pos_id, pos_loc = pos
+            pos_loc = (int(pos_loc[0]), int(pos_loc[1]))
+            cv2.circle(image, center=pos_loc, radius=radius, color=(0, 255, 0), thickness=-1)
+            # rr, cc = skimage.draw.circle(int(pos_loc[1]), int(pos_loc[0]), radius)
+            # image[rr, cc] = [0, 255, 0]
+
+        ## draw sketch
+        # joint id (0 - r ankle, 1 - r knee, 2 - r hip, 3 - l hip, 4 - l knee,
+        #           5 - l ankle, 6 - pelvis, 7 - thorax, 8 - upper neck,
+        #           9 - head top, 10 - r wrist, 11 - r elbow, 12 - r shoulder,
+        #           13 - l shoulder, 14 - l elbow, 15 - l wrist)
+        #
+        #               9
+        #               8
+        #         12 ** 7 ** 13
+        #        *      *      *
+        #       11      *       14
+        #      *        *         *
+        #     10    2 * 6 * 3     15
+        #           *       *
+        #           1       4
+        #           *       *
+        #           0       5
+
+        lines = [
+            [(0, 1), [100, 255, 100]],
+            [(1, 2), [50, 255, 50]],
+            [(2, 6), [0, 255, 0]],  # right leg
+            [(3, 4), [100, 100, 255]],
+            [(4, 5), [50, 50, 255]],
+            [(6, 3), [0, 0, 255]],  # left leg
+            [(6, 7), [255, 255, 100]],
+            [(7, 8), [255, 150, 50]],  # body
+            [(8, 9), [255, 200, 100]],  # head
+            [(10, 11), [255, 100, 255]],
+            [(11, 12), [255, 50, 255]],
+            [(12, 8), [255, 0, 255]],  # right hand
+            [(8, 13), [0, 255, 255]],
+            [(13, 14), [100, 255, 255]],
+            [(14, 15), [200, 255, 255]]  # left hand
+        ]
+        for line in lines:
+            start, end = line[0]
+            if (start in joint_pos) and (end in joint_pos):
+                cv2.line(
+                    image,
+                    (int(joint_pos[start][0]), int(joint_pos[start][1])),
+                    (int(joint_pos[end][0]), int(joint_pos[end][1])),  # up-left and botton-right
+                    line[1],
+                    thick)
+                # rr, cc, val = skimage.draw.line_aa(int(joint_pos[start][1]), int(joint_pos[start][0]), int(joint_pos[end][1]), int(joint_pos[end][0]))
+                # image[rr, cc] = line[1]
+
+    if save_name is not None:
+        # cv2.imwrite(save_name, image)
+        save_image(image, save_name)
+    return image
+
+
 def frame(I=None, second=5, saveable=True, name='frame', cmap=None, fig_idx=12836):
     """Display a frame(image). Make sure OpenAI Gym render() is disable before using it.
 
