Exemplo n.º 1
0
def main():

    use_gpu = False

    if use_gpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    else:
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

    input_img = tf.placeholder(tf.float32, shape=[1, None, None, 3])

    # _1, _2, cpm, paf = light_openpose(input_img, is_training=False)
    cpm, paf = lightweight_openpose(input_img, num_pafs=26, num_joints=14, is_training=False)
    saver = tf.train.Saver()

    total_img = 0
    total_time = 0
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        saver.restore(sess, params['test_model'])
        print('#---------Successfully loaded trained model.---------#')
        save(sess)
        if params['img_path'] is not None:
            img_ori = cv2.imread(params['img_path'])
            print(np.shape(img_ori))
            img_data = cv2.cvtColor(img_ori, code=cv2.COLOR_BGR2RGB)
            img_data = cv2.resize(img_data, (256, 256))
            img = img_data / 255.

            start_time = time.time()
            heatmap, _paf = sess.run([cpm, paf], feed_dict={input_img: [img]})
            end_time = time.time()
            print(heatmap.shape)
            print(_paf.shape)
            canvas, joint_list, person_to_joint_assoc, joints = decode_pose(img_data, params, heatmap[0], _paf[0])
            decode_time = time.time()
            print('inference + decode time == {}'.format(decode_time - start_time))

            total_img += 1
            total_time += (end_time - start_time)
            canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
            cv2_imshow(canvas)
        else:
            print('Nothing to process.')
Exemplo n.º 2
0
    def predict(self, img_original):

        #preprocess image to get 'model_input'
        model_input = self.preprocess(img_original)

        # execute model inference
        result = self.model.execute([model_input])

        # postprocessing: use the heatmaps (the second output of model) to get the joins and limbs for human body
        # Note: the model has multiple outputs, here we used a simplified method, which only uses heatmap for body joints
        #       and the heatmap has shape of [1,14], each value correspond to the position of one of the 14 joints.
        #       The value is the index in the 92*92 heatmap (flatten to one dimension)
        heatmaps = result[1]

        # calculate the scale of original image over heatmap, Note: image_original.shape[0] is height
        scale = np.array([
            img_original.shape[1] / heatmap_width,
            img_original.shape[0] / heatmap_height
        ])

        canvas, joint_list = decode_pose(heatmaps[0], scale, img_original)
        return canvas, joint_list
Exemplo n.º 3
0
def main():

    use_gpu = False

    if use_gpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    else:
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

    input_img = tf.placeholder(tf.float32, shape=[1, None, None, 3])

    # _1, _2, cpm, paf = light_openpose(input_img, is_training=False)
    cpm, paf = lightweight_openpose(input_img, num_pafs=26, num_joints=14, is_training=False)
    saver = tf.train.Saver()

    total_img = 0
    total_time = 0
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        saver.restore(sess, params['test_model'])
        print('#---------Successfully loaded trained model.---------#')
        if 'video_path'in params.keys() and params['video_path'] is not None:
            # video_capture = cv2.VideoCapture('rtsp://*****:*****@10.24.1.238')
            video_capture = cv2.VideoCapture(params['video_path'])
            fps = video_capture.get(cv2.CAP_PROP_FPS)
            start_second = 0
            start_frame = fps * start_second
            video_capture.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
            while True:
                retval, img_data = video_capture.read()
                if not retval:
                    break
                img_data = cv2.cvtColor(img_data, code=cv2.COLOR_BGR2RGB)
                orih, oriw, c = img_data.shape
                img = cv2.resize(img_data, (256, 256)) / 255.
                start_time = time.time()
                heatmap, _paf = sess.run([cpm, paf], feed_dict={input_img: [img]})
                end_time = time.time()

                canvas, joint_list, person_to_joint_assoc = decode_pose(img, params, heatmap[0], _paf[0])
                decode_time = time.time()
                print ('inference + decode time == {}'.format(decode_time - start_time))
                total_img += 1
                total_time += (end_time-start_time)
                canvas = canvas.astype(np.float32)
                canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
                cv2.imshow('result', canvas)
                cv2.waitKey(1)
        elif params['img_path'] is not None:
            for img_name in os.listdir(params['img_path']):
                if img_name.split('.')[-1] != 'jpg':
                    continue
                img_data = cv2.imread(os.path.join(params['img_path'], img_name))
                img_data = cv2.cvtColor(img_data, code=cv2.COLOR_BGR2RGB)
                img = cv2.resize(img_data, (256, 256)) / 255.
                start_time = time.time()
                heatmap, _paf = sess.run([cpm, paf], feed_dict={input_img: [img]})
                end_time = time.time()
                canvas, joint_list, person_to_joint_assoc = decode_pose(img, params, heatmap[0], _paf[0])
                canvas = canvas.astype(np.float32)
                canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
                decode_time = time.time()
                print('inference + decode time == {}'.format(decode_time - start_time))
                total_img += 1
                total_time += (end_time - start_time)
                cv2.imshow('result', canvas)
                cv2.waitKey(0)

        else:
            print('Nothing to process.')
Exemplo n.º 4
0
def main():

    use_gpu = params['gpu']

    if use_gpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    else:
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

    with open(params['json_file'], encoding='utf-8') as fp:
        labels = json.load(fp)

    input_img = tf.placeholder(tf.float32, shape=[1, None, None, 3])

    cpm, paf = lightweight_openpose(input_img,
                                    num_pafs=26,
                                    num_joints=14,
                                    is_training=False)

    saver = tf.train.Saver()
    total_img = 0

    predictions = []

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        saver.restore(sess, params['test_model'])
        print('#---------Successfully loaded trained model.---------#')

        for label in labels:

            img_id = label['image_id']

            img_ori = cv2.imread(
                os.path.join(params['img_path'], img_id + '.jpg'))
            img_data = cv2.cvtColor(img_ori, code=cv2.COLOR_BGR2RGB)
            img_data = cv2.resize(img_data,
                                  (params['input_size'], params['input_size']))
            img = img_data / 255.

            heatmap, _paf = sess.run([cpm, paf], feed_dict={input_img: [img]})
            canvas, joint_list, person_to_joint_assoc, joints = decode_pose(
                img_data, params, heatmap[0], _paf[0])

            predict = label
            predict['image_id'] = img_id
            kps = {}
            human = 1
            factorY = img_ori.shape[0] / img_data.shape[0]
            factorX = img_ori.shape[1] / img_data.shape[1]

            for joint in joints:
                for i in range(14):
                    joint[3 * i] *= factorX
                    joint[3 * i + 1] *= factorY
                if type(joint) == type([]):
                    kps['human' + str(human)] = joint
                else:
                    kps['human' + str(human)] = joint.tolist()
                human += 1

            # print (human)
            # print (kps)
            # for person, joint in kps.items():
            #     assert len(joint) == 14 * 3
            #     for i in range(14):
            #         x = int(joint[i*3])
            #         y = int(joint[i*3+1])
            #         cv2.circle(img_ori, (x, y), 3, (255, 255, 255), thickness=-1)
            #         cv2.putText(img_ori, str(i), (x, y),cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 1)
            # cv2.imshow('test', img_ori)
            # cv2.waitKey(0)

            predict['keypoint_annotations'] = kps
            predictions.append(predict)

            # break
            total_img += 1
            print('processing number: {}'.format(total_img))
            # break
    with open(params['save_json'], 'w') as fw:
        json.dump(predictions, fw)