Exemple #1
0
def detection(image):
    time_start=time.time()
    image_batch = data_to_input(image)

    # Compute prediction with the CNN
    outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
    scmap, locref, pairwise_diff = predict.extract_cnn_output(outputs_np, cfg, dataset.pairwise_stats)

    detections = extract_detections(cfg, scmap, locref, pairwise_diff)
    unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)
    person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)

    time_end=time.time()
    print('totally cost', time_end-time_start)
    return person_conf_multi
Exemple #2
0
def run_predict(frame, sess, inputs, outputs, cfg, dataset, sm, draw_multi):

    tf.reset_default_graph()
    image= frame
    image_batch = data_to_input(frame)

    # Compute prediction_n with the CNN
    outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
    scmap, locref, pairwise_diff = predict.extract_cnn_output(outputs_np, cfg, dataset.pairwise_stats)
    detections = extract_detections(cfg, scmap, locref, pairwise_diff)
    unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)
    m=time.time()
    person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)
    img = np.copy(image)
    visim_multi = img.copy()
    draw_multi.draw(visim_multi, dataset, person_conf_multi, image)
    return pos_array.round().astype(int)
Exemple #3
0
    image = video.get_frame(i / video.fps)

    ##########
    ## By pose-tensorflow

    image_batch = data_to_input(image)

    # Compute prediction with the CNN
    outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
    scmap, locref, pairwise_diff = predict.extract_cnn_output(
        outputs_np, cfg, dataset.pairwise_stats)

    detections = extract_detections(cfg, scmap, locref, pairwise_diff)
    unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(
        sm, detections)
    person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array,
                                                 pos_array)

    #####

    # Add library to draw image
    image_img = Image.fromarray(image)

    # Prepare saving image with points of pose
    draw = ImageDraw.Draw(image_img)

    #####

    people_num = 0
    people_real_num = 0

    people_num = person_conf_multi.size / (point_num * 2)
def test_net(visualise, cache_scoremaps, development):
    logging.basicConfig(level=logging.INFO)

    cfg = load_config()
    dataset = create_dataset(cfg)
    dataset.set_shuffle(False)

    sm = SpatialModel(cfg)
    sm.load()

    draw_multi = PersonDraw()

    from_cache = "cached_scoremaps" in cfg
    if not from_cache:
        sess, inputs, outputs = setup_pose_prediction(cfg)

    if cache_scoremaps:
        out_dir = cfg.scoremap_dir
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

    pairwise_stats = dataset.pairwise_stats
    num_images = dataset.num_images if not development else min(
        10, dataset.num_images)
    coco_results = []

    for k in range(num_images):
        print('processing image {}/{}'.format(k, num_images - 1))

        batch = dataset.next_batch()

        cache_name = "{}.mat".format(batch[Batch.data_item].coco_id)

        if not from_cache:
            outputs_np = sess.run(outputs,
                                  feed_dict={inputs: batch[Batch.inputs]})
            scmap, locref, pairwise_diff = extract_cnn_output(
                outputs_np, cfg, pairwise_stats)

            if cache_scoremaps:
                if visualise:
                    img = np.squeeze(batch[Batch.inputs]).astype('uint8')
                    pose = argmax_pose_predict(scmap, locref, cfg.stride)
                    arrows = argmax_arrows_predict(scmap, locref,
                                                   pairwise_diff, cfg.stride)
                    visualize.show_arrows(cfg, img, pose, arrows)
                    visualize.waitforbuttonpress()
                    continue

                out_fn = os.path.join(out_dir, cache_name)
                dict = {
                    'scoremaps': scmap.astype('float32'),
                    'locreg_pred': locref.astype('float32'),
                    'pairwise_diff': pairwise_diff.astype('float32')
                }
                scipy.io.savemat(out_fn, mdict=dict)
                continue
        else:
            # cache_name = '1.mat'
            full_fn = os.path.join(cfg.cached_scoremaps, cache_name)
            mlab = scipy.io.loadmat(full_fn)
            scmap = mlab["scoremaps"]
            locref = mlab["locreg_pred"]
            pairwise_diff = mlab["pairwise_diff"]

        detections = extract_detections(cfg, scmap, locref, pairwise_diff)
        unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(
            sm, detections)
        person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array,
                                                     pos_array)

        if visualise:
            img = np.squeeze(batch[Batch.inputs]).astype('uint8')
            # visualize.show_heatmaps(cfg, img, scmap, pose)
            """
            # visualize part detections after NMS
            visim_dets = visualize_detections(cfg, img, detections)
            plt.imshow(visim_dets)
            plt.show()
            visualize.waitforbuttonpress()
            """

            #            """
            visim_multi = img.copy()
            draw_multi.draw(visim_multi, dataset, person_conf_multi)

            plt.imshow(visim_multi)
            plt.show()
            visualize.waitforbuttonpress()
        #            """

        if cfg.use_gt_segm:
            coco_img_results = pose_predict_with_gt_segm(
                scmap, locref, cfg.stride, batch[Batch.data_item].gt_segm,
                batch[Batch.data_item].coco_id)
            coco_results += coco_img_results
            if len(coco_img_results):
                dataset.visualize_coco(coco_img_results,
                                       batch[Batch.data_item].visibilities)

    if cfg.use_gt_segm:
        with open('predictions_with_segm.json', 'w') as outfile:
            json.dump(coco_results, outfile)

    sess.close()
def return_pose(image, image2, keeper, referee):

    keeper_id = -1
    ref_ids = []

    team_class_num = 0.25  # sys.maxsize
    ref_num = 0.10

    contours, mask = get_contours(image)

    all_info = []

    image_batch = data_to_input(image)

    # Compute prediction with the CNN
    outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
    scmap, locref, pairwise_diff = predict.extract_cnn_output(
        outputs_np, cfg, dataset.pairwise_stats)

    detections = extract_detections(cfg, scmap, locref, pairwise_diff)
    unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(
        sm, detections)
    person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array,
                                                 pos_array)

    cl = [(255, 0, 0), (0, 0, 255), (0, 255, 0), (255, 255, 0), (0, 255, 255),
          (255, 0, 255), (0, 0, 0)]

    names = {
        6: 'rightShoulder',
        5: 'leftShoulder',
        11: 'leftHip',
        12: 'rightHip',
        14: 'rightKnee',
        13: 'leftKnee',
        16: 'rightAnkle',
        15: 'leftAnkle'
    }

    allTeamClassificationFeatures_upper = []
    allTeamClassificationFeatures_lower = []

    # TODO - Use mid-back to knee features
    for i in range(len(person_conf_multi)):

        x = 0
        player_parts = {}
        min_x = image2.shape[0] - 1
        min_y = image2.shape[1] - 1

        for j in range(17):
            printFlag = True
            if j in [5, 6, 15, 16, 11, 12, 13, 14]:

                ptsy, ptsx = person_conf_multi[i, j, :]
                # print(ptsx, ptsy, ':', min_x, min_y)

                if ptsy > 0.0 and ptsx > 0.0:
                    player_parts.update({names[j]: {'x': ptsx, 'y': ptsy}})

                    x = x + 1
                    if min_x > ptsx:
                        min_x = ptsx
                        min_y = ptsy

        # print('\n', min_x, min_y)

        teamClassificationFeatures_upper = []
        teamClassificationFeatures_lower = []

        allXCoords_upper = []
        allYCoords_upper = []
        allXCoords_lower = []
        allYCoords_lower = []

        leftUpperPointFound = False
        rightUpperPointFound = False
        leftMidPointFound = False
        rightMidPointFound = False
        leftLowerPointFound = False
        rightLowerPointFound = False

        if 'rightShoulder' in list(player_parts.keys()):
            allXCoords_upper.append(int(player_parts['rightShoulder']['x']))
            allYCoords_upper.append(int(player_parts['rightShoulder']['y']))
            rightUpperPointFound = True
        if 'leftShoulder' in list(player_parts.keys()):
            allXCoords_upper.append(int(player_parts['leftShoulder']['x']))
            allYCoords_upper.append(int(player_parts['leftShoulder']['y']))
            leftUpperPointFound = True

        if 'rightHip' in list(player_parts.keys()):
            allXCoords_upper.append(int(player_parts['rightHip']['x']))
            allYCoords_upper.append(int(player_parts['rightHip']['y']))
            allXCoords_lower.append(int(player_parts['rightHip']['x']))
            allYCoords_lower.append(int(player_parts['rightHip']['y']))
            rightMidpointFound = True
        if 'leftHip' in list(player_parts.keys()):
            allXCoords_upper.append(int(player_parts['leftHip']['x']))
            allYCoords_upper.append(int(player_parts['leftHip']['y']))
            allXCoords_lower.append(int(player_parts['leftHip']['x']))
            allYCoords_lower.append(int(player_parts['leftHip']['y']))
            leftMidPointFound = True

        if 'rightKnee' in list(player_parts.keys()):
            allXCoords_lower.append(int(player_parts['rightKnee']['x']))
            allYCoords_lower.append(int(player_parts['rightKnee']['y']))
            rightLowerPointFound = True
        if 'leftKnee' in list(player_parts.keys()):
            allXCoords_lower.append(int(player_parts['leftKnee']['x']))
            allYCoords_lower.append(int(player_parts['leftKnee']['y']))
            leftLowerPointFound = True

        allXCoords_upper.sort()
        allYCoords_upper.sort()
        allXCoords_lower.sort()
        allYCoords_lower.sort()

        if (len(allXCoords_upper) < 3) or (len(allXCoords_lower) < 3):
            continue

        if mask[int(min_x), int(min_y)] == 0:
            continue

        # Upper body features
        if len(allXCoords_upper) == 4:
            if 0 > allXCoords_upper[0]:
                allXCoords_upper[0] = 0
            if 0 > allXCoords_upper[3]:
                allXCoords_upper[3] = 0
            if 0 > allXCoords_upper[0] and 0 > allXCoords_upper[3]:
                allXCoords_upper[3] = 2
            if 0 > allYCoords_upper[0]:
                allYCoords_upper[0] = 0
            if 0 > allYCoords_upper[3]:
                allYCoords_upper[3] = 0
            if 0 > allYCoords_upper[0] and 0 > allYCoords_upper[3]:
                allYCoords_upper[3] = 2
            allXCoords_upper.sort()
            allYCoords_upper.sort()
            allColors = image2[allXCoords_upper[0]:allXCoords_upper[3],
                               allYCoords_upper[0]:allYCoords_upper[3]]
            image = cv2.rectangle(
                image, (allYCoords_upper[0], allXCoords_upper[0]),
                (int(allYCoords_upper[3] -
                     abs(allYCoords_upper[3] - allYCoords_upper[0]) / 2),
                 allXCoords_upper[3]), (255, 0, 0), 1)
            chans = cv2.split(allColors)
            colors = ("b", "g", "r")
            for (chan, color) in zip(chans, colors):
                hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
                teamClassificationFeatures_upper.extend(
                    [np.int64(x[0]) for x in hist])
                teamClassificationFeatures_upper = [
                    float(i) / max(teamClassificationFeatures_upper)
                    for i in teamClassificationFeatures_upper
                ]
            allTeamClassificationFeatures_upper.append(
                teamClassificationFeatures_upper)

        if len(allXCoords_upper) == 3:
            if 0 > allXCoords_upper[0]:
                allXCoords_upper[0] = 0
            if 0 > allXCoords_upper[2]:
                allXCoords_upper[2] = 0
            if 0 > allXCoords_upper[0] and 0 > allXCoords_upper[2]:
                allXCoords_upper[2] = 10
            if 0 > allYCoords_upper[0]:
                allYCoords_upper[0] = 0
            if 0 > allYCoords_upper[2]:
                allYCoords_upper[2] = 0
            if 0 > allYCoords_upper[0] and 0 > allYCoords_upper[2]:
                allYCoords_upper[2] = 10
            allXCoords_upper.sort()
            allYCoords_upper.sort()
            allColors = image2[allXCoords_upper[0]:allXCoords_upper[2],
                               allYCoords_upper[0]:allYCoords_upper[2]]
            image = cv2.rectangle(
                image, (allYCoords_upper[0], allXCoords_upper[0]),
                (int(allYCoords_upper[3] -
                     abs(allYCoords_upper[2] - allYCoords_upper[0]) / 2),
                 allXCoords_upper[2]), (255, 0, 0), 1)
            chans = cv2.split(allColors)
            colors = ("b", "g", "r")
            for (chan, color) in zip(chans, colors):
                hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
                teamClassificationFeatures_upper.extend(
                    [np.int64(x[0]) for x in hist])
                teamClassificationFeatures_upper = [
                    float(i) / max(teamClassificationFeatures_upper)
                    for i in teamClassificationFeatures_upper
                ]
            allTeamClassificationFeatures_upper.append(
                teamClassificationFeatures_upper)

        if len(allXCoords_upper) == 2:
            # if (leftUpperPointFound and rightMidPointFound) or (leftMidPointFound and rightUpperPointFound):
            if 0 > allXCoords_upper[0]:
                allXCoords_upper[0] = 0
            if 0 > allXCoords_upper[1]:
                allXCoords_upper[1] = 0
            if 0 > allXCoords_upper[0] and 0 > allXCoords_upper[1]:
                allXCoords_upper[1] = 2
            if 0 > allYCoords_upper[0]:
                allYCoords_upper[0] = 0
            if 0 > allYCoords_upper[1]:
                allYCoords_upper[1] = 0
            if 0 > allYCoords_upper[0] and 0 > allYCoords_upper[1]:
                allYCoords_upper[1] = 2
            allXCoords_upper.sort()
            allYCoords_upper.sort()
            allColors = image2[allXCoords_upper[0]:allXCoords_upper[1],
                               allYCoords_upper[0]:allYCoords_upper[1]]
            image = cv2.rectangle(
                image, (allYCoords_upper[0], allXCoords_upper[0]),
                (int(allYCoords_upper[1] -
                     abs(allYCoords_upper[3] - allYCoords_upper[0]) / 2),
                 allXCoords_upper[1]), (255, 0, 0), 1)
            '''reqInLinePoints = bresenham_march(image2, [allXCoords[0], allYCoords[0]], [allXCoords[1], allYCoords[1]])
			allColors = []
			for point in reqInLinePoints:
				allColors.append(point[1].reshape(1,3))
			reqInLinePoints = bresenham_march(image2, [allXCoords[0], allYCoords[1]], [allXCoords[1], allYCoords[0]])
			for point in reqInLinePoints:
				allColors.append(point[1].reshape(1,3))'''
            chans = cv2.split(allColors)
            colors = ("b", "g", "r")
            for (chan, color) in zip(chans, colors):
                hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
                teamClassificationFeatures_upper.extend(
                    [np.int64(x[0]) for x in hist])
                teamClassificationFeatures_upper = [
                    float(i) / max(teamClassificationFeatures_upper)
                    for i in teamClassificationFeatures_upper
                ]
            allTeamClassificationFeatures_upper.append(
                teamClassificationFeatures_upper)

        # Lower body features
        if len(allXCoords_lower) == 4:
            if 0 > allXCoords_lower[0]:
                allXCoords_lower[0] = 0
            if 0 > allXCoords_lower[3]:
                allXCoords_lower[3] = 0
            if 0 > allXCoords_lower[0] and 0 > allXCoords_lower[3]:
                allXCoords_lower[3] = 2
            if 0 > allYCoords_lower[0]:
                allYCoords_lower[0] = 0
            if 0 > allYCoords_lower[3]:
                allYCoords_lower[3] = 0
            if 0 > allYCoords_lower[0] and 0 > allYCoords_lower[3]:
                allYCoords_lower[3] = 2
            allXCoords_lower.sort()
            allYCoords_lower.sort()
            allColors = image2[allXCoords_lower[0]:allXCoords_lower[3],
                               allYCoords_lower[0]:allYCoords_lower[3]]
            image = cv2.rectangle(image,
                                  (allYCoords_lower[0], allXCoords_lower[0]),
                                  (allYCoords_lower[3], allXCoords_lower[3]),
                                  (255, 0, 0), 1)
            chans = cv2.split(allColors)
            colors = ("b", "g", "r")
            for (chan, color) in zip(chans, colors):
                hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
                teamClassificationFeatures_lower.extend(
                    [np.int64(x[0]) for x in hist])
                teamClassificationFeatures_lower = [
                    float(i) / max(teamClassificationFeatures_lower)
                    for i in teamClassificationFeatures_lower
                ]
            allTeamClassificationFeatures_lower.append(
                teamClassificationFeatures_lower)

        if len(allXCoords_lower) == 3:
            if 0 > allXCoords_lower[0]:
                allXCoords_lower[0] = 0
            if 0 > allXCoords_lower[2]:
                allXCoords_lower[2] = 0
            if 0 > allXCoords_lower[0] and 0 > allXCoords_lower[2]:
                allXCoords_lower[2] = 10
            if 0 > allYCoords_lower[0]:
                allYCoords_lower[0] = 0
            if 0 > allYCoords_lower[2]:
                allYCoords_lower[2] = 0
            if 0 > allYCoords_lower[0] and 0 > allYCoords_lower[2]:
                allYCoords_lower[2] = 10
            allXCoords_lower.sort()
            allYCoords_lower.sort()
            allColors = image2[allXCoords_lower[0]:allXCoords_lower[2],
                               allYCoords_lower[0]:allYCoords_lower[2]]
            image = cv2.rectangle(image,
                                  (allYCoords_lower[0], allXCoords_lower[0]),
                                  (allYCoords_lower[2], allXCoords_lower[2]),
                                  (255, 0, 0), 1)
            chans = cv2.split(allColors)
            colors = ("b", "g", "r")
            for (chan, color) in zip(chans, colors):
                hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
                teamClassificationFeatures_lower.extend(
                    [np.int64(x[0]) for x in hist])
                teamClassificationFeatures_lower = [
                    float(i) / max(teamClassificationFeatures_lower)
                    for i in teamClassificationFeatures_lower
                ]
            allTeamClassificationFeatures_lower.append(
                teamClassificationFeatures_lower)

        if len(allXCoords_lower) == 2:
            # if (leftMidPointFound and rightLowerPointFound) or (leftLowerPointFound and rightMidPointFound):
            if 0 > allXCoords_lower[0]:
                allXCoords_lower[0] = 0
            if 0 > allXCoords_lower[1]:
                allXCoords_lower[1] = 0
            if 0 > allXCoords_lower[0] and 0 > allXCoords_lower[1]:
                allXCoords_lower[1] = 2
            if 0 > allYCoords_lower[0]:
                allYCoords_lower[0] = 0
            if 0 > allYCoords_lower[1]:
                allYCoords_lower[1] = 0
            if 0 > allYCoords_lower[0] and 0 > allYCoords_lower[1]:
                allYCoords_lower[1] = 2
            allXCoords_lower.sort()
            allYCoords_lower.sort()
            allColors = image2[allXCoords_lower[0]:allXCoords_lower[1],
                               allYCoords_lower[0]:allYCoords_lower[1]]
            image = cv2.rectangle(image,
                                  (allYCoords_lower[0], allXCoords_lower[0]),
                                  (allYCoords_lower[1], allXCoords_lower[1]),
                                  (255, 0, 0), 1)
            '''reqInLinePoints = bresenham_march(image2, [allXCoords[0], allYCoords[0]], [allXCoords[1], allYCoords[1]])
			allColors = []
			for point in reqInLinePoints:
				allColors.append(point[1].reshape(1,3))
			reqInLinePoints = bresenham_march(image2, [allXCoords[0], allYCoords[1]], [allXCoords[1], allYCoords[0]])
			for point in reqInLinePoints:
				allColors.append(point[1].reshape(1,3))'''
            chans = cv2.split(allColors)
            colors = ("b", "g", "r")
            for (chan, color) in zip(chans, colors):
                hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
                teamClassificationFeatures_lower.extend(
                    [np.int64(x[0]) for x in hist])
                teamClassificationFeatures_lower = [
                    float(i) / max(teamClassificationFeatures_lower)
                    for i in teamClassificationFeatures_lower
                ]
            allTeamClassificationFeatures_lower.append(
                teamClassificationFeatures_lower)
        '''distance_keep = np.sqrt(np.sum(np.square(np.array(teamClassificationFeatures) - np.array(keeper))))
		if distance_keep < team_class_num:
			team_class_num = distance_keep
			keeper_id = i
		
		distance_ref = np.sqrt(np.sum(np.square(np.array(teamClassificationFeatures) - np.array(referee))))
		if distance_ref < ref_num:
			ref_ids.append(i)'''

        temp = None

        all_info.append([i, temp, player_parts, [min_x, min_y]])

    teamClassifier_upper = KMeans(n_clusters=2)
    teamLabels_upper = teamClassifier_upper.fit_predict(
        allTeamClassificationFeatures_upper)
    teamClassifierDBSCAN_upper = DBSCAN(
        eps=0.5, min_samples=2,
        metric='euclidean').fit(allTeamClassificationFeatures_upper)
    teamLabelsDBSCAN_upper = teamClassifierDBSCAN_upper.labels_

    teamClassifier_lower = KMeans(n_clusters=2)
    teamLabels_lower = teamClassifier_lower.fit_predict(
        allTeamClassificationFeatures_lower)
    teamClassifierDBSCAN_lower = DBSCAN(
        eps=0.5, min_samples=2,
        metric='euclidean').fit(allTeamClassificationFeatures_lower)
    teamLabelsDBSCAN_lower = teamClassifierDBSCAN_lower.labels_

    isKeeperFound = False
    isRefFound = False

    # Mapping clusters
    # Create cost matrix
    CM = np.zeros(shape=(2, 2))
    for i in range(2):
        for j in range(2):
            CM[i][j] =  calc_dist_between_points(teamClassifier_upper.cluster_centers_[i], \
             teamClassifier_lower.cluster_centers_[j], \
             len(teamClassifier_lower.cluster_centers_))

    # Hungarian algorithm
    row_ind, col_ind = linear_sum_assignment(CM)
    upper_to_lower_map = dict(zip(row_ind, col_ind))
    lower_to_upper_map = dict(zip(col_ind, row_ind))

    for player_itr in range(len(all_info)):

        dist_from_centroid_upper = np.sqrt(np.sum(np.square(np.array(allTeamClassificationFeatures_upper[player_itr]) - \
         np.array(teamClassifier_upper.cluster_centers_[teamLabels_upper[player_itr]]))))
        dist_from_centroid_lower = np.sqrt(np.sum(np.square(np.array(allTeamClassificationFeatures_lower[player_itr]) - \
         np.array(teamClassifier_lower.cluster_centers_[teamLabels_lower[player_itr]]))))

        if (teamLabelsDBSCAN_upper[player_itr]
                == -1) or (teamLabelsDBSCAN_lower[player_itr] == -1):
            all_info[player_itr][1] = -1
        else:
            all_info[player_itr][1] = teamLabels_upper[
                player_itr] * 10 + teamLabels_lower[player_itr]
        '''elif (upper_to_lower_map[teamLabels_upper[player_itr]] == teamLabels_lower[player_itr]):
			all_info[player_itr][1] = teamLabels_upper[player_itr]
		else:
			if dist_from_centroid_upper < dist_from_centroid_lower:
				all_info[player_itr][1] = teamLabels_upper[player_itr]
			else:
				all_info[player_itr][1] = lower_to_upper_map[teamLabels_lower[player_itr]]'''
        '''if all_info[player_itr][0] == keeper_id:
			isKeeperFound = True
			all_info[player_itr][1] = 2'''
        '''if all_info[player_itr][1] in ref_ids:
			isRefFound = True
			all_info[player_itr][1] = 3'''

    print("ref:", ref_ids, "keeper:", keeper_id, "score", team_class_num)

    return all_info, isKeeperFound, isRefFound, image
# Read image from file
file_name_ext = sys.argv[1] ## example: test_single_03.png
file_name = file_name_ext.split('.')[0] ## example: test_single_03
file_name_input = 'testset/' + file_name_ext
image = imread(file_name_input, mode='RGB')

image_batch = data_to_input(image)

# Compute prediction with the CNN
outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
scmap, locref, pairwise_diff = predict.extract_cnn_output(outputs_np, cfg, dataset.pairwise_stats)

detections = extract_detections(cfg, scmap, locref, pairwise_diff)
unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)
person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)

print('person_conf_multi: ')
print(type(person_conf_multi))
print(person_conf_multi)

# img = np.copy(image)

# visim_multi = img.copy()

# fig = plt.imshow(visim_multi)
# draw_multi.draw(visim_multi, dataset, person_conf_multi)
# fig.axes.get_xaxis().set_visible(False)
# fig.axes.get_yaxis().set_visible(False)
#
# plt.show()
Exemple #7
0
def return_pose(image , image2 , keeper , referee):
	
	keeper_id = -1
	ref_ids= []
	
	team_class_num = 0.25 # sys.maxsize
	ref_num = 0.10
	
	contours, mask = get_contours(image)
		
	all_info = []

	image_batch = data_to_input(image)

	# Compute prediction with the CNN
	outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
	scmap, locref, pairwise_diff = predict.extract_cnn_output(outputs_np, cfg, dataset.pairwise_stats)

	detections = extract_detections(cfg, scmap, locref, pairwise_diff)
	unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)
	person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)

	cl = [(255,0,0),(0,0,255),(0,255,0),(255,255,0),(0,255,255),(255,0,255),(0,0,0)]

	names = {6:'rightShoulder', 5:'leftShoulder', 11:'leftHip', 12:'rightHip',14:'rightKnee',13:'leftKnee',16:'rightAnkle',15:'leftAnkle'}	

	allTeamClassificationFeatures = []
	allTeamClassificationFeaturesDBSCAN = []
	
	# TODO - Use mid-back to knee features
	for i in range(len(person_conf_multi)):
		
		x = 0
		player_parts = {}
		min_x = image2.shape[0] - 1
		min_y = image2.shape[1] - 1

		for j in range(17):
			
			printFlag = True
			if j in [5, 6, 15, 16, 11, 12, 13, 14]:
				ptsy, ptsx = person_conf_multi[i,j,:]

				if ptsy > 0.0 and ptsx > 0.0:
					player_parts.update({names[j] : {'x' : ptsx , 'y' :ptsy} })
					x = x + 1
					
					if min_x > ptsx:
						min_x = ptsx
						min_y = ptsy
		    
		teamClassificationFeatures = []
		teamClassificationFeaturesDBSCAN = []
		
		allXCoords = []
		allYCoords = []

		leftUpperPointFound = False
		rightUpperPointFound = False
		leftLowerPointFound = False
		rightLowerPointFound = False
		hipPointsFound = False
		
		if 'rightShoulder' in list(player_parts.keys()):
			allXCoords.append(int(player_parts['rightShoulder']['x']))
			allYCoords.append(int(player_parts['rightShoulder']['y']))
			rightUpperPointFound = True
		if 'leftShoulder' in list(player_parts.keys()):
			allXCoords.append(int(player_parts['leftShoulder']['x']))
			allYCoords.append(int(player_parts['leftShoulder']['y']))
			leftUpperPointFound = True

		if 'rightKnee' in list(player_parts.keys()):
			allXCoords.append(int(player_parts['rightKnee']['x']))
			allYCoords.append(int(player_parts['rightKnee']['y']))
			rightLowerPointFound = True		
		if 'leftKnee' in list(player_parts.keys()):
			allXCoords.append(int(player_parts['leftKnee']['x']))
			allYCoords.append(int(player_parts['leftKnee']['y']))
			leftLowerPointFound = True

		allXCoords.sort()
		allYCoords.sort()
			
		if len(allXCoords) < 3:
			continue

		hipPointsFound = 'rightHip' in list(player_parts.keys()) and 'leftHip' in list(player_parts.keys())
		if hipPointsFound:
			hipPointsFound &= abs(int(player_parts['leftHip']['y']) - int(player_parts['rightHip']['y'])) > 5
		if hipPointsFound:
			allHipXCoords = []
			allHipYCoords = []
			allHipXCoords.append(int(player_parts['rightHip']['x']))
			allHipYCoords.append(int(player_parts['rightHip']['y']))
			allHipXCoords.append(int(player_parts['leftHip']['x'])) 
			allHipYCoords.append(int(player_parts['leftHip']['y']))
			allHipXCoords.sort()
			allHipYCoords.sort()
		
		if mask[int(min_x) , int(min_y)] == 0:
			continue

		if len(allXCoords) == 4:
			if 0 > allXCoords[0]:
				allXCoords[0] = 0
			if 0 > allXCoords[3]:
				allXCoords[3] = 0
			if 0 > allXCoords[0] and 0 > allXCoords[3]:
				allXCoords[3] = 2
			if 0 > allYCoords[0]:
				allYCoords[0] = 0
			if 0 > allYCoords[3]:
				allYCoords[3] = 0
			if 0 > allYCoords[0] and 0 > allYCoords[3]:
				allYCoords[3] = 2 
			allXCoords.sort()
			allYCoords.sort() 
			newXLower = int(allXCoords[0]+abs(allXCoords[3]-allXCoords[0])*0.4)
			newXUpper = int(allXCoords[3]-abs(allXCoords[3]-allXCoords[0])*0.2)
			if hipPointsFound:
				newYLower = int(allHipYCoords[0])
				newYUpper = int(allHipYCoords[-1])
			else:
				newYLower = int(allYCoords[0])
				newYUpper = int(allYCoords[3]) 
			allColors = image2[newXLower:newXUpper, newYLower:newYUpper]
			allColorsDBSCAN = image2[allXCoords[0]:allXCoords[-1], allYCoords[0]:allYCoords[-1]]
			image = cv2.rectangle(image, (newYLower, newXLower), (newYUpper, newXUpper), (255, 0, 0), 1)
			chans = cv2.split(allColors)
			colors = ("b", "g", "r")
			for (chan, color) in zip(chans, colors):
				hist = cv2.calcHist([chan], [0], None, [64], [0, 256])
				teamClassificationFeatures.extend([np.int64(x[0]) for x in hist])
				teamClassificationFeatures = [float(i)/max(teamClassificationFeatures) for i in teamClassificationFeatures]
			allTeamClassificationFeatures.append(teamClassificationFeatures)
			chans = cv2.split(allColorsDBSCAN)
			colors = ("b", "g", "r")
			for (chan, color) in zip(chans, colors):
				hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
				teamClassificationFeaturesDBSCAN.extend([np.int64(x[0]) for x in hist])
				teamClassificationFeaturesDBSCAN = [float(i)/max(teamClassificationFeaturesDBSCAN) for i in teamClassificationFeaturesDBSCAN]
			allTeamClassificationFeaturesDBSCAN.append(teamClassificationFeaturesDBSCAN)

		if len(allXCoords) == 3:
			if 0 > allXCoords[0]:
				allXCoords[0] = 0
			if 0 > allXCoords[2]:
				allXCoords[2] = 0
			if 0 > allXCoords[0] and 0 > allXCoords[2]:
				allXCoords[2] = 10 
			if 0 > allYCoords[0]:
				allYCoords[0] = 0
			if 0 > allYCoords[2]:
				allYCoords[2] = 0
			if 0 > allYCoords[0] and 0 > allYCoords[2]:
				allYCoords[2] = 10  
			allXCoords.sort()
			allYCoords.sort()
			newXLower = int(allXCoords[0]+abs(allXCoords[2]-allXCoords[0])*0.4)
			newXUpper = int(allXCoords[2]-abs(allXCoords[2]-allXCoords[0])*0.2)
			if hipPointsFound:
				newYLower = int(allHipYCoords[0])
				newYUpper = int(allHipYCoords[-1])
			else:
				newYLower = int(allYCoords[0])
				newYUpper = int(allYCoords[2])
			allColors = image2[newXLower:newXUpper, newYLower:newYUpper]
			allColorsDBSCAN = image2[allXCoords[0]:allXCoords[-1], allYCoords[0]:allYCoords[-1]]
			image = cv2.rectangle(image, (newYLower, newXLower), (newYUpper, newXUpper), (255, 0, 0), 1)
			chans = cv2.split(allColors)
			colors = ("b", "g", "r")
			for (chan, color) in zip(chans, colors):
				hist = cv2.calcHist([chan], [0], None, [64], [0, 256])
				teamClassificationFeatures.extend([np.int64(x[0]) for x in hist])
				teamClassificationFeatures = [float(i)/max(teamClassificationFeatures) for i in teamClassificationFeatures]
			allTeamClassificationFeatures.append(teamClassificationFeatures)
			chans = cv2.split(allColorsDBSCAN)
			colors = ("b", "g", "r")
			for (chan, color) in zip(chans, colors):
				hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
				teamClassificationFeaturesDBSCAN.extend([np.int64(x[0]) for x in hist])
				teamClassificationFeaturesDBSCAN = [float(i)/max(teamClassificationFeaturesDBSCAN) for i in teamClassificationFeaturesDBSCAN]
			allTeamClassificationFeaturesDBSCAN.append(teamClassificationFeaturesDBSCAN)

		if len(allXCoords) == 2:
			# if (leftUpperPointFound and rightLowerPointFound) or (leftLowerPointFound and rightUpperPointFound):
			if 0 > allXCoords[0]:
				allXCoords[0] = 0
			if 0 > allXCoords[1]:
				allXCoords[1] = 0
			if 0 > allXCoords[0] and 0 > allXCoords[1]:
				allXCoords[1] = 2 
			if 0 > allYCoords[0]:
				allYCoords[0] = 0
			if 0 > allYCoords[1]:
				allYCoords[1] = 0
			if 0 > allYCoords[0] and 0 > allYCoords[1]:
				allYCoords[1] = 2
			allXCoords.sort()
			allYCoords.sort()
			newXLower = int(allXCoords[0]+abs(allXCoords[1]-allXCoords[0])*0.35)
			newXUpper = int(allXCoords[1]*0.9) 
			allColors = image2[newXLower:newXUpper, allYCoords[0]:allYCoords[1]]
			allColorsDBSCAN = image2[allXCoords[0]:allXCoords[1], allYCoords[0]:allYCoords[1]]
			image = cv2.rectangle(image, (allYCoords[0], int(allXCoords[0]+abs(allXCoords[1]-allXCoords[0])*0.35)), (allYCoords[1], allXCoords[1]), (255, 0, 0), 1)
			chans = cv2.split(allColors)
			colors = ("b", "g", "r")
			for (chan, color) in zip(chans, colors):
				hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
				teamClassificationFeatures.extend([np.int64(x[0]) for x in hist])
				teamClassificationFeatures = [float(i)/max(teamClassificationFeatures) for i in teamClassificationFeatures]
			allTeamClassificationFeatures.append(teamClassificationFeatures)
			chans = cv2.split(allColorsDBSCAN)
			colors = ("b", "g", "r")
			for (chan, color) in zip(chans, colors):
				hist = cv2.calcHist([chan], [0], None, [3], [0, 256])
				teamClassificationFeaturesDBSCAN.extend([np.int64(x[0]) for x in hist])
				teamClassificationFeaturesDBSCAN = [float(i)/max(teamClassificationFeaturesDBSCAN) for i in teamClassificationFeaturesDBSCAN]
			allTeamClassificationFeaturesDBSCAN.append(teamClassificationFeaturesDBSCAN)
		
		temp = None
		
		all_info.append([i, temp, player_parts, [min_x,min_y]])
	
	teamClassifier = KMeans(n_clusters = 2)
	teamLabels = teamClassifier.fit_predict(allTeamClassificationFeatures)
	teamClassifierDBSCAN = DBSCAN(eps=0.5, min_samples=2, metric='euclidean').fit(allTeamClassificationFeaturesDBSCAN)
	teamLabelsDBSCAN = teamClassifierDBSCAN.labels_

	isKeeperFound = False
	isRefFound = False

	dist_from_centroids = []
	for player_itr in range(len(all_info)):

		dist_from_centroid = np.sqrt(np.sum(np.square(np.array(allTeamClassificationFeatures[player_itr]) - \
			np.array(teamClassifier.cluster_centers_[teamLabels[player_itr]]))))
		dist_from_centroids.append(dist_from_centroid)

	normed_dist_from_centroids = [float(i)/sum(dist_from_centroids) for i in dist_from_centroids]
	print('normes dist', normed_dist_from_centroids)
		
	for player_itr in range(len(all_info)):

		if (teamLabelsDBSCAN[player_itr] == -1):
			all_info[player_itr][1] = -1
		else:
			all_info[player_itr][1] = teamLabels[player_itr]
			
	print("ref:", ref_ids,"keeper:", keeper_id, "score", team_class_num)	
		
	return all_info, isKeeperFound, isRefFound, image
Exemple #8
0
def process_body_image(req):
    process_width = int(rospy.get_param("~process_width", 320))
    process_height = int(rospy.get_param("~process_height", 240))

    try:
        cv_image = bridge.imgmsg_to_cv2(req.image, "bgr8")
    except CvBridgeError as e:
        rospy.logerr(e)
    height, width, channels = cv_image.shape
    resize_rio_width = 1.0 * width / process_width
    resize_rio_height = 1.0 * height / process_height

    resized_cv_image = cv2.resize(cv_image, (process_width, process_height))
    image_batch = data_to_input(resized_cv_image)
    outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
    scmap, locref, pairwise_diff = predict.extract_cnn_output(
        outputs_np, cfg, dataset.pairwise_stats)
    detections = extract_detections(cfg, scmap, locref, pairwise_diff)
    unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(
        sm, detections)
    person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array,
                                                 pos_array)
    num_people = person_conf_multi.shape[0]
    conf_min_count = rospy.get_param("conf_min_count", 5)  # 至少三个点才被认为能够接受
    bodys_response = BodyPoseResponse()
    for pidx in range(num_people):
        if np.sum(person_conf_multi[pidx, :, 0] > 0) < conf_min_count:
            continue
        body_rect = person_conf_multi[pidx].tolist()
        body_rect_filterd = list(
            filter(lambda x: x[0] > 0 and x[1] > 0, body_rect))
        if len(body_rect_filterd) < conf_min_count:
            continue
        body_info = BodyInfo()
        body_info.nose = [
            body_rect[0][0] * resize_rio_width,
            body_rect[0][1] * resize_rio_height
        ]
        body_info.right_eye = [
            body_rect[1][0] * resize_rio_width,
            body_rect[1][1] * resize_rio_height
        ]
        body_info.left_eye = [
            body_rect[2][0] * resize_rio_width,
            body_rect[2][1] * resize_rio_height
        ]
        body_info.right_ear = [
            body_rect[3][0] * resize_rio_width,
            body_rect[3][1] * resize_rio_height
        ]
        body_info.left_ear = [
            body_rect[4][0] * resize_rio_width,
            body_rect[4][1] * resize_rio_height
        ]
        body_info.right_arm_top = [
            body_rect[5][0] * resize_rio_width,
            body_rect[5][1] * resize_rio_height
        ]
        body_info.left_arm_top = [
            body_rect[6][0] * resize_rio_width,
            body_rect[6][1] * resize_rio_height
        ]
        body_info.right_arm_middle = [
            body_rect[7][0] * resize_rio_width,
            body_rect[7][1] * resize_rio_height
        ]
        body_info.left_arm_middle = [
            body_rect[8][0] * resize_rio_width,
            body_rect[8][1] * resize_rio_height
        ]
        body_info.right_arm_bottom = [
            body_rect[9][0] * resize_rio_width,
            body_rect[9][1] * resize_rio_height
        ]
        body_info.left_arm_bottom = [
            body_rect[10][0] * resize_rio_width,
            body_rect[10][1] * resize_rio_height
        ]
        body_info.right_leg_top = [
            body_rect[11][0] * resize_rio_width,
            body_rect[11][1] * resize_rio_height
        ]
        body_info.left_leg_top = [
            body_rect[12][0] * resize_rio_width,
            body_rect[12][1] * resize_rio_height
        ]
        body_info.right_leg_middle = [
            body_rect[13][0] * resize_rio_width,
            body_rect[13][1] * resize_rio_height
        ]
        body_info.left_leg_middle = [
            body_rect[14][0] * resize_rio_width,
            body_rect[14][1] * resize_rio_height
        ]
        body_info.right_leg_bottom = [
            body_rect[15][0] * resize_rio_width,
            body_rect[15][1] * resize_rio_height
        ]
        body_info.left_leg_bottom = [
            body_rect[16][0] * resize_rio_width,
            body_rect[16][1] * resize_rio_height
        ]
        bodys_response.body_poses.append(body_info)
    return bodys_response
Exemple #9
0
def main():
    start_time=time.time()
    print("main hai")
    tf.reset_default_graph()
    cfg = load_config("demo/pose_cfg_multi.yaml")
    dataset = create_dataset(cfg)
    sm = SpatialModel(cfg)
    sm.load()
    draw_multi = PersonDraw()
    # Load and setup CNN part detector
    sess, inputs, outputs = predict.setup_pose_prediction(cfg)

    # Read image from file
    dir=os.listdir("stick")
    k=0
    cap=cv2.VideoCapture(0)
    i=0
    while (cap.isOpened()):
            if i%20 == 0:                   
                ret, orig_frame= cap.read()
                if ret==True:
                    frame = cv2.resize(orig_frame, (0, 0), fx=0.30, fy=0.30)
                    image= frame
                    sse=0
                    mse=0
                    
                    image_batch = data_to_input(frame)

                    # Compute prediction with the CNN
                    outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})

                    scmap, locref, pairwise_diff = predict.extract_cnn_output(outputs_np, cfg, dataset.pairwise_stats)

                    detections = extract_detections(cfg, scmap, locref, pairwise_diff)

                    unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)

                    person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)
                    img = np.copy(image)
                    #coor = PersonDraw.draw()
                    visim_multi = img.copy()
                    co1=draw_multi.draw(visim_multi, dataset, person_conf_multi)
                    plt.imshow(visim_multi)
                    plt.show()
                    visualize.waitforbuttonpress()
                    #print("this is draw : ", co1)
                    if k==1:
                        qwr = np.zeros((1920,1080,3), np.uint8)

                        cv2.line(qwr, co1[5][0], co1[5][1],(255,0,0),3)
                        cv2.line(qwr, co1[7][0], co1[7][1],(255,0,0),3)
                        cv2.line(qwr, co1[6][0], co1[6][1],(255,0,0),3)
                        cv2.line(qwr, co1[4][0], co1[4][1],(255,0,0),3)

                        cv2.line(qwr, co1[9][0], co1[9][1],(255,0,0),3)
                        cv2.line(qwr, co1[11][0], co1[11][1],(255,0,0),3)
                        cv2.line(qwr, co1[8][0], co1[8][1],(255,0,0),3)
                        cv2.line(qwr, co1[10][0], co1[10][1],(255,0,0),3)
                        # In[9]:
                        cv2.imshow('r',qwr)
                        qwr2="stick/frame"+str(k)+".jpg"
                        qw1 = cv2.cvtColor(qwr, cv2.COLOR_BGR2GRAY)
                        qw2= cv2.cvtColor(qwr2, cv2.COLOR_BGR2GRAY)

                        fig = plt.figure("Images")
                        images = ("Original", qw1), ("Contrast", qw2)
                        for (i, (name, image)) in enumerate(images):
                                ax = fig.add_subplot(1, 3, i + 1)
                                ax.set_title(name)
                        plt.imshow(hash(tuple(image)))
                        # compare the images
                        s,m=compare_images(qw1, qw2, "Image1 vs Image2")
                        k+=1
                        sse=s
                        mse=m

                else:
                    break
    elapsed= time.time()-start_time
    #print("sse score : ", sse)
    print("Mean squared error : ", elapsed/100)
    cap.release()
    cv2.destroyAllWindows()
def video2poseframe(video_name):
    import numpy as np

    sys.path.append(os.path.dirname(__file__) + "/../")

    from scipy.misc import imread, imsave

    from config import load_config
    from dataset.factory import create as create_dataset
    from nnet import predict
    from util import visualize
    from dataset.pose_dataset import data_to_input

    from multiperson.detections import extract_detections
    from multiperson.predict import SpatialModel, eval_graph, get_person_conf_multicut
    from multiperson.visualize import PersonDraw, visualize_detections

    import matplotlib.pyplot as plt

    from PIL import Image, ImageDraw

    import random

    cfg = load_config("demo/pose_cfg_multi.yaml")

    dataset = create_dataset(cfg)

    sm = SpatialModel(cfg)
    sm.load()

    # Load and setup CNN part detector
    sess, inputs, outputs = predict.setup_pose_prediction(cfg)

    ################

    video = read_video(video_name)

    video_frame_number = int(video.duration * video.fps) ## duration: second / fps: frame per second
    video_frame_ciphers = math.ceil(math.log(video_frame_number, 10)) ## ex. 720 -> 3

    if not os.path.exists('testset/' + video_name):
        os.makedirs('testset/' + video_name)

    for i in range(0, video_frame_number):
        image = video.get_frame(i/video.fps)

        ######################

        image_batch = data_to_input(image)

        # Compute prediction with the CNN
        outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
        scmap, locref, pairwise_diff = predict.extract_cnn_output(outputs_np, cfg, dataset.pairwise_stats)

        detections = extract_detections(cfg, scmap, locref, pairwise_diff)
        unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)
        person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)

        print('person_conf_multi: ')
        print(type(person_conf_multi))
        print(person_conf_multi)

        # Add library to save image
        image_img = Image.fromarray(image)

        # Save image with points of pose
        draw = ImageDraw.Draw(image_img)

        people_num = 0
        point_num = 17
        print('person_conf_multi.size: ')
        print(person_conf_multi.size)
        people_num = person_conf_multi.size / (point_num * 2)
        people_num = int(people_num)
        print('people_num: ')
        print(people_num)

        point_i = 0 # index of points
        point_r = 5 # radius of points

        people_real_num = 0
        for people_i in range(0, people_num):
            point_color_r = random.randrange(0, 256)
            point_color_g = random.randrange(0, 256)
            point_color_b = random.randrange(0, 256)
            point_color = (point_color_r, point_color_g, point_color_b, 255)
            point_count = 0
            for point_i in range(0, point_num):
                if person_conf_multi[people_i][point_i][0] + person_conf_multi[people_i][point_i][1] != 0: # If coordinates of point is (0, 0) == meaningless data
                    point_count = point_count + 1
            if point_count > 5: # If there are more than 5 point in person, we define he/she is REAL PERSON
                people_real_num = people_real_num + 1
                for point_i in range(0, point_num):
                    draw.ellipse((person_conf_multi[people_i][point_i][0] - point_r, person_conf_multi[people_i][point_i][1] - point_r, person_conf_multi[people_i][point_i][0] + point_r, person_conf_multi[people_i][point_i][1] + point_r), fill=point_color)

        print('people_real_num: ')
        print(people_real_num)

        video_name_result = 'testset/' + video_name + '/frame_pose_' + str(i).zfill(video_frame_ciphers) + '.jpg'
        image_img.save(video_name_result, "JPG")
def video2posevideo(video_name):
    time_start = time.clock()

    import numpy as np

    sys.path.append(os.path.dirname(__file__) + "/../")

    from scipy.misc import imread, imsave

    from config import load_config
    from dataset.factory import create as create_dataset
    from nnet import predict
    from util import visualize
    from dataset.pose_dataset import data_to_input

    from multiperson.detections import extract_detections
    from multiperson.predict import SpatialModel, eval_graph, get_person_conf_multicut
    from multiperson.visualize import PersonDraw, visualize_detections

    import matplotlib.pyplot as plt

    from PIL import Image, ImageDraw, ImageFont
    font = ImageFont.truetype("./font/NotoSans-Bold.ttf", 24)

    import random

    cfg = load_config("demo/pose_cfg_multi.yaml")

    dataset = create_dataset(cfg)

    sm = SpatialModel(cfg)
    sm.load()

    draw_multi = PersonDraw()

    # Load and setup CNN part detector
    sess, inputs, outputs = predict.setup_pose_prediction(cfg)

    ################

    video = read_video(video_name)

    video_frame_number = int(video.duration * video.fps) ## duration: second / fps: frame per second
    video_frame_ciphers = math.ceil(math.log(video_frame_number, 10)) ## ex. 720 -> 3

    pose_frame_list = []

    point_r = 3 # radius of points
    point_min = 10 # threshold of points - If there are more than point_min points in person, we define he/she is REAL PERSON
    part_min = 3 # threshold of parts - If there are more than part_min parts in person, we define he/she is REAL PERSON / part means head, arm and leg
    point_num = 17 # There are 17 points in 1 person

    def ellipse_set(person_conf_multi, people_i, point_i):
        return (person_conf_multi[people_i][point_i][0] - point_r, person_conf_multi[people_i][point_i][1] - point_r, person_conf_multi[people_i][point_i][0] + point_r, person_conf_multi[people_i][point_i][1] + point_r)

    def line_set(person_conf_multi, people_i, point_i, point_j):
        return (person_conf_multi[people_i][point_i][0], person_conf_multi[people_i][point_i][1], person_conf_multi[people_i][point_j][0], person_conf_multi[people_i][point_j][1])

    def draw_ellipse_and_line(draw, person_conf_multi, people_i, a, b, c, point_color):
        draw.ellipse(ellipse_set(person_conf_multi, people_i, a), fill=point_color)
        draw.ellipse(ellipse_set(person_conf_multi, people_i, b), fill=point_color)
        draw.ellipse(ellipse_set(person_conf_multi, people_i, c), fill=point_color)
        draw.line(line_set(person_conf_multi, people_i, a, b), fill=point_color, width=5)
        draw.line(line_set(person_conf_multi, people_i, b, c), fill=point_color, width=5)

    for i in range(0, video_frame_number):
        image = video.get_frame(i/video.fps)

        ######################

        image_batch = data_to_input(image)

        # Compute prediction with the CNN
        outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
        scmap, locref, pairwise_diff = predict.extract_cnn_output(outputs_np, cfg, dataset.pairwise_stats)

        detections = extract_detections(cfg, scmap, locref, pairwise_diff)
        unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)
        person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)

        # print('person_conf_multi: ')
        # print(type(person_conf_multi))
        # print(person_conf_multi)

        # Add library to save image
        image_img = Image.fromarray(image)

        # Save image with points of pose
        draw = ImageDraw.Draw(image_img)

        people_num = 0
        people_real_num = 0
        people_part_num = 0

        people_num = person_conf_multi.size / (point_num * 2)
        people_num = int(people_num)
        print('people_num: ' + str(people_num))

        for people_i in range(0, people_num):
            point_color_r = random.randrange(0, 256)
            point_color_g = random.randrange(0, 256)
            point_color_b = random.randrange(0, 256)
            point_color = (point_color_r, point_color_g, point_color_b, 255)
            point_list = []
            point_count = 0
            point_i = 0 # index of points
            part_count = 0 # count of parts in THAT person

            # To find rectangle which include that people - list of points x, y coordinates
            people_x = []
            people_y = []

            for point_i in range(0, point_num):
                if person_conf_multi[people_i][point_i][0] + person_conf_multi[people_i][point_i][1] != 0: # If coordinates of point is (0, 0) == meaningless data
                    point_count = point_count + 1
                    point_list.append(point_i)

            # Draw each parts
            if (5 in point_list) and (7 in point_list) and (9 in point_list): # Draw left arm
                draw_ellipse_and_line(draw, person_conf_multi, people_i, 5, 7, 9, point_color)
                part_count = part_count + 1
            if (6 in point_list) and (8 in point_list) and (10 in point_list): # Draw right arm
                draw_ellipse_and_line(draw, person_conf_multi, people_i, 6, 8, 10, point_color)
                part_count = part_count + 1
            if (11 in point_list) and (13 in point_list) and (15 in point_list): # Draw left leg
                draw_ellipse_and_line(draw, person_conf_multi, people_i, 11, 13, 15, point_color)
                part_count = part_count + 1
            if (12 in point_list) and (14 in point_list) and (16 in point_list): # Draw right leg
                draw_ellipse_and_line(draw, person_conf_multi, people_i, 12, 14, 16, point_color)
                part_count = part_count + 1
            if point_count >= point_min:
                people_real_num = people_real_num + 1
                for point_i in range(0, point_num):
                    if person_conf_multi[people_i][point_i][0] + person_conf_multi[people_i][point_i][1] != 0: # If coordinates of point is (0, 0) == meaningless data
                        draw.ellipse(ellipse_set(person_conf_multi, people_i, point_i), fill=point_color)
                        people_x.append(person_conf_multi[people_i][point_i][0])
                        people_y.append(person_conf_multi[people_i][point_i][1])
                # Draw rectangle which include that people
                draw.rectangle([min(people_x), min(people_y), max(people_x), max(people_y)], fill=point_color, outline=5)


            if part_count >= part_min:
                people_part_num = people_part_num + 1

        draw.text((0, 0), 'People(by point): ' + str(people_real_num) + ' (threshold = ' + str(point_min) + ')', (0,0,0), font=font)
        draw.text((0, 32), 'People(by line): ' + str(people_part_num) + ' (threshold = ' + str(part_min) + ')', (0,0,0), font=font)
        draw.text((0, 64), 'Frame: ' + str(i) + '/' + str(video_frame_number), (0,0,0), font=font)
        draw.text((0, 96), 'Total time required: ' + str(round(time.clock() - time_start, 1)) + 'sec', (0,0,0))

        print('people_real_num: ' + str(people_real_num))
        print('people_part_num: ' + str(people_part_num))
        print('frame: ' + str(i))

        image_img_numpy = np.asarray(image_img)

        pose_frame_list.append(image_img_numpy)

    video_pose = ImageSequenceClip(pose_frame_list, fps=video.fps)
    video_pose.write_videofile("testset/" + video_name + "_pose.mp4", fps=video.fps)

    print("Time(s): " + str(time.clock() - time_start))
def test_net(visualise, cache_scoremaps, development):
    logging.basicConfig(level=logging.INFO)

    cfg = load_config()
    dataset = create_dataset(cfg)
    dataset.set_shuffle(False)

    sm = SpatialModel(cfg)
    sm.load()

    draw_multi = PersonDraw()

    from_cache = "cached_scoremaps" in cfg
    if not from_cache:
        sess, inputs, outputs = setup_pose_prediction(cfg)

    if cache_scoremaps:
        out_dir = cfg.scoremap_dir
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

    pairwise_stats = dataset.pairwise_stats
    num_images = dataset.num_images if not development else min(10, dataset.num_images)
    coco_results = []

    for k in range(num_images):
        print('processing image {}/{}'.format(k, num_images-1))

        batch = dataset.next_batch()

        cache_name = "{}.mat".format(batch[Batch.data_item].coco_id)

        if not from_cache:
            outputs_np = sess.run(outputs, feed_dict={inputs: batch[Batch.inputs]})
            scmap, locref, pairwise_diff = extract_cnn_output(outputs_np, cfg, pairwise_stats)

            if cache_scoremaps:
                if visualise:
                    img = np.squeeze(batch[Batch.inputs]).astype('uint8')
                    pose = argmax_pose_predict(scmap, locref, cfg.stride)
                    arrows = argmax_arrows_predict(scmap, locref, pairwise_diff, cfg.stride)
                    visualize.show_arrows(cfg, img, pose, arrows)
                    visualize.waitforbuttonpress()
                    continue

                out_fn = os.path.join(out_dir, cache_name)
                dict = {'scoremaps': scmap.astype('float32'),
                        'locreg_pred': locref.astype('float32'),
                        'pairwise_diff': pairwise_diff.astype('float32')}
                scipy.io.savemat(out_fn, mdict=dict)
                continue
        else:
            #cache_name = '1.mat'
            full_fn = os.path.join(cfg.cached_scoremaps, cache_name)
            mlab = scipy.io.loadmat(full_fn)
            scmap = mlab["scoremaps"]
            locref = mlab["locreg_pred"]
            pairwise_diff = mlab["pairwise_diff"]

        detections = extract_detections(cfg, scmap, locref, pairwise_diff)
        unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)
        person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)

        if visualise:
            img = np.squeeze(batch[Batch.inputs]).astype('uint8')
            #visualize.show_heatmaps(cfg, img, scmap, pose)

            """
            # visualize part detections after NMS
            visim_dets = visualize_detections(cfg, img, detections)
            plt.imshow(visim_dets)
            plt.show()
            visualize.waitforbuttonpress()
            """

#            """
            visim_multi = img.copy()
            draw_multi.draw(visim_multi, dataset, person_conf_multi)

            plt.imshow(visim_multi)
            plt.show()
            visualize.waitforbuttonpress()
#            """


        if cfg.use_gt_segm:
            coco_img_results = pose_predict_with_gt_segm(scmap, locref, cfg.stride, batch[Batch.data_item].gt_segm,
                                                      batch[Batch.data_item].coco_id)
            coco_results += coco_img_results
            if len(coco_img_results):
                dataset.visualize_coco(coco_img_results, batch[Batch.data_item].visibilities)

    if cfg.use_gt_segm:
        with open('predictions_with_segm.json', 'w') as outfile:
            json.dump(coco_results, outfile)

    sess.close()