-
Notifications
You must be signed in to change notification settings - Fork 31
/
video.py
110 lines (87 loc) · 3.5 KB
/
video.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
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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Patrick Wieschollek <mail@patwie.com>
"""
Just a wrapper to ease the process of handling videos in OpenCV.
"""
import os
import cv2
import numpy as np
class ReadFrameException(Exception):
pass
class Reader(object):
"""Read Videos"""
def __init__(self, video_path):
super(Reader, self).__init__()
assert os.path.isfile(video_path), "The video %s does not exists!" % video_path
self.video_path = video_path
self.vid = cv2.VideoCapture(video_path)
# print("read video {}\n fps: {}\n frames: {}\n shape: {}".format(video_path, self.fps,
# self.frames, self.shape))
self.frames_cache = None
@property
def fps(self):
return self.vid.get(cv2.CAP_PROP_FPS)
@property
def frames(self):
if self.frames_cache:
return self.frames_cache
self.frames_cache = int(self.vid.get(cv2.CAP_PROP_FRAME_COUNT))
return self.frames_cache
@property
def pos(self):
return int(self.vid.get(cv2.CAP_PROP_POS_FRAMES))
@property
def height(self):
return int(self.vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
@property
def width(self):
return int(self.vid.get(cv2.CAP_PROP_FRAME_WIDTH))
@property
def total_seconds(self):
return self.frames / self.fps
@property
def shape(self):
w = int(self.vid.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(self.vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
return (h, w)
def read(self, scale=None):
success, frame = self.vid.read()
if not success:
raise ReadFrameException('Cannot read frame %i from video %s with msg %s' %
(self.pos, self.video_path, success))
if scale:
frame = cv2.resize(frame, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
return frame
def jump(self, frame=None, ms=None):
assert frame is not ms, "Use either frame or ms, not both!"
if frame:
if frame >= self.frames:
raise ReadFrameException('Cannot jump to frame (frame does not exists)')
self.vid.set(cv2.CAP_PROP_POS_FRAMES, frame)
if ms:
self.vid.set(cv2.CAP_PROP_POS_MSEC, ms)
# print("jumped to frame %i" % self.vid.get(cv2.CAP_PROP_POS_FRAMES))
class Writer(object):
"""Write Videos"""
def __init__(self, fn, height, width, channels=3, fps=30):
super(Writer, self).__init__()
self.width = width
self.height = height
self.channels = channels
self.fn = fn
print("writing to %s with shape (%i, %i, %i)" % (fn, height, width, channels))
cc4 = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
self.out = cv2.VideoWriter(fn, cc4, fps, (width, height), isColor=(channels == 3))
def write(self, frame):
if len(frame.shape) == 2:
frame = frame[:, :, None]
assert frame.shape == (self.height, self.width, self.channels), "{} vs {}".format(frame.shape,
(self.width, self.height,
self.channels))
frame = np.round(frame)
frame = np.clip(frame, 0, 255)
frame = frame.astype('uint8')
self.out.write(frame)
def finish(self):
self.out.release()