在视频中加入标注——Demo

发表于

有些数据以视频的形式给出，在观察数据集的时候就需要对其进行处理，将对应的标注信息显示在视频中。

示例代码

这段代码依次读取视频文件，对其进行标注之后显示出来。

Main.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import numpy as np
import cv2
import h5py
import time
import os
from drawJoints import DrawJoints

CONNECTED_PAIRS=[[10,11],[11,12],[13,14],[14,15],[0,1],[1,2],[3,4],[4,5]]


def showVid(cap, pred,title='frame'):
    idx = 0
    fps = cap.get(cv2.CAP_PROP_FPS)
    while (True):
        # t1=time.time()
        if idx >= pred.shape[0]:
            return
        ret, frame = cap.read()
        frame = DrawJoints(frame, pred[idx, ::], CONNECTED_PAIRS)
        cv2.imshow(title, frame)
        idx = idx + 1
        if cv2.waitKey(int(1000 / fps)) & 0xFF == ord(' '):
            while True:
                if (cv2.waitKey(1) & 0xFF == ord('q')):
                    return
                if (cv2.waitKey(1) & 0xFF == ord(' ')):
                    break
        if ( cap.get(cv2.CAP_PROP_FRAME_COUNT) <= cap.get(cv2.CAP_PROP_POS_FRAMES) ):
            return
        if (cv2.waitKey(2) & 0xFF == ord('q')) :
            return

def getFileList(data_dir,pred_dir,idx):
    '''
    Assume idx in ['S9','S11']
    :param data_dir:
    :param pred_dir:
    :param idx:
    :return:
    '''
    V=[] #Video
    P=[] #Predict
    video_files=[i for i in os.listdir(os.path.join(data_dir,idx,'Videos')) if (i.split('.')[-1]=='mp4' and (not i[0]=='_') )]
    pred_files= [i for i in os.listdir(os.path.join(pred_dir,idx,'StackedHourglassFineTuned240')) if i.split('.')[-1]=='h5']
    video_files.sort()
    pred_files = [p.replace('_',' ') for p in pred_files]
    pred_files.sort()
    pred_files = [p.replace(' ', '_') for p in pred_files]

    ip=0
    for i in range(len(video_files)):
        cur_file=video_files[i]
        if idx=='S11' and cur_file=='Directions.54138969.mp4':
            continue
        else:
            V.append(os.path.join(data_dir,idx,'Videos',cur_file))
            P.append(os.path.join(os.path.join(pred_dir,idx,'StackedHourglassFineTuned240'),pred_files[ip]))
            ip=ip+1
    return V,P

def main():
    vidDir='/data3/Human36M/raw_data'
    predDir='/data2/guoyu/workspace/3d-pose-baseline/temporal/data/h36m'
    for S in ['S9','S11']:
        V,P = getFileList(vidDir,predDir,S)
        for i in range(len(V)):
            pred_name=P[i]
            filename=V[i]
            pred_data = h5py.File(pred_name)
            pred = np.array(pred_data['poses'], dtype=np.int)
            cap = cv2.VideoCapture(filename)
            # print int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),'\t',int(cap.get(cv2.CAP_PROP_FRAME_COUNT)- pred.shape[0]),'\t' ,S,filename.split('/')[-1].split('.')[0]
            # cap.release()
            #
            fps = cap.get(cv2.CAP_PROP_FPS)
            showVid(cap, pred,title=filename.split('/')[-1])
            cap.release()
            cv2.destroyAllWindows()


if __name__ == '__main__':
    main()

使用opencv的VideoCapture函数可以方便地读取视频数据，可以通过cap的get函数来获得对应的视频属性，例如帧率、帧数、长宽等信息。
cv2.waitKey(2) & 0xFF == ord('q')语句可以判断键盘的输入。

drawJoints.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import os
import numpy as np
import cv2

JoName='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'
JoName=JoName.split(',')

def DrawJoints(I,joints,CONNECTED_PAIRS=None,Jo=JoName):
    '''
    Assume joints Dim : P x 2
    :param I:
    :param joints:
    :param CONNECTED_PAIRS:
    :return:
    '''
    txt='Missing Joints:\n'
    # Background rectangle for text
    cv2.rectangle(I, (25, 25), (300, 250), (0, 0, 0), -1)
    font = cv2.FONT_HERSHEY_SIMPLEX
    for i in range(joints.shape[0]):
        if joints[i,0]<5 and joints[i,1]<5 :
            txt=txt+ str(Jo[i]) +'\n'
        cv2.circle(I, (joints[i,0], joints[i,1]), 2, (0, 0, 255), -1)
    if CONNECTED_PAIRS:
        I=drawLine(I,joints,CONNECTED_PAIRS)

    # show text
    y0, dy = 50, 25
    for i, txt in enumerate(txt.split('\n')):
        y = y0 + i * dy
        cv2.putText(I, txt, (50, y), font, .8, (0, 255, 0), 1, 2)
    # print(txt)
    return I

def drawLine(I,joints,CONNECTED_PAIRS):
    N=len(CONNECTED_PAIRS)
    for i in range(N):
        p1=joints[CONNECTED_PAIRS[i][0]]
        p2=joints[CONNECTED_PAIRS[i][1]]
        p1=totuple(p1)
        p2 = totuple(p2)
        cv2.line(I,p1,p2,(0,255,0),2)
    return I

def totuple(a):
    try:
        return tuple(totuple(i) for i in a)
    except TypeError:
        return a

putText 函数用于在图像上输出文字，需要注意的是对于字符串中的换行符\n，会显示乱码，需要另外处理。

相关文章