TempEC/labelling.py at master · B-Step62/TempEC · GitHub

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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import cv2
import sys
import csv
import os.path
import numpy as np

def player(video_path):
    #動画のパス
    root, video_name = os.path.split(video_path)

    #csvのパス
    video_root, ext = os.path.splitext(video_name)
    csv_path = '{0}.csv'.format(video_root)


    #動画読み込み
    cap = cv2.VideoCapture(video_path)

    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    #labels初期化(csvがあれば読み込み)
    if os.path.exists(csv_path):
        f = open(csv_path, 'r')
        reader = csv.reader(f)
        labels = []
        for row in reader:
            for i in row:
                labels.append(int(float(i)))
    else:
        labels = [0 for i in range(total)]


    frame_num = 0
    block_num = 0

    while(block_num * 100 < total):
        #フレームを100枚ずつ配列に格納
        frames = []
        while(cap.isOpened()):

            ret, frame = cap.read()

            if frame is None:
                break

            frames.append(frame)

            sys.stdout.write("\r%d" % len(frames))
            sys.stdout.flush()

            if len(frames) >= 100:
                break

        #100枚について繰り返し
        while frame_num - block_num * 100 < len(frames):
            frame = frames[frame_num - block_num * 100]

            frame = cv2.resize(frame,(1280, 720))

            cv2.putText(frame,str(frame_num + 1) + '/' + str(total),(15,45),cv2.FONT_HERSHEY_PLAIN, 3, (255,255,255),2,cv2.LINE_AA)

            cv2.imshow("Frame", frame)

            key = cv2.waitKey()

            if key == 121:    #y
                labels[frame_num] = 1
                frame_num += 1

            elif key == 110:     #n
                labels[frame_num] = 0
                frame_num += 1

            elif key == 98:      #b
                if frame_num - block_num * 100 > 0:
                    frame_num -= 1

            elif key == 115:     #s
                frame_num = block_num * 100 + len(frames)

            elif key == 32:      #space
                cap.release()
                cv2.destroyAllWindows()
                save(labels,csv_path)
                sys.exit()

            #else:
                #labels[frame_num] = 0
                #frame_num += 1

            key = 0

        save(labels,csv_path)

        block_num += 1


    cap.release()
    cv2.destroyAllWindows()

    save(labels,csv_path)

    sys.exit()


def save(labels,csv_path):
    np.savetxt(csv_path,labels,delimiter=',')


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print('not enough parms')
        exit()

    video_path = sys.argv[1]

    player(video_path)