/
videostreamer.py
99 lines (95 loc) · 3.59 KB
/
videostreamer.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
# -*- coding: utf-8 -*-
"""
Created on Wed Nov 27 15:55:49 2019
@author: achuiko
"""
import cv2
import glob
import os
class VideoStreamer(object):
""" Class to help process image streams. Three types of possible inputs:"
1.) USB Webcam.
2.) A directory of images (files in directory matching 'img_glob').
3.) A video file, such as an .mp4 or .avi file.
"""
def __init__(self, basedir, camid, height, width, skip, img_glob):
self.cap = []
self.camera = False
self.video_file = False
self.listing = []
self.sizer = [height, width]
self.i = 0
self.skip = skip
self.maxlen = 1000000
# If the "basedir" string is the word camera, then use a webcam.
if basedir == "camera/" or basedir == "camera":
print('==> Processing Webcam Input.')
self.cap = cv2.VideoCapture(camid)
self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
self.camera = True
else:
# Try to open as a video.
self.cap = cv2.VideoCapture(basedir)
lastbit = basedir[-4:len(basedir)]
if (type(self.cap) == list or not self.cap.isOpened()) and (lastbit == '.mp4'):
raise IOError('Cannot open movie file')
elif type(self.cap) != list and self.cap.isOpened() and (lastbit != '.txt'):
print('==> Processing Video Input.')
num_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
self.listing = range(0, num_frames)
self.listing = self.listing[::self.skip]
self.camera = True
self.video_file = True
self.maxlen = len(self.listing)
else:
print('==> Processing Image Directory Input.')
search = os.path.join(basedir, img_glob)
self.listing = glob.glob(search)
self.listing.sort()
self.listing = self.listing[::self.skip]
self.maxlen = len(self.listing)
if self.maxlen == 0:
raise IOError('No images were found (maybe bad \'--img_glob\' parameter?)')
def read_image(self, impath, img_size):
""" Read image as grayscale and resize to img_size.
Inputs
impath: Path to input image.
img_size: (W, H) tuple specifying resize size.
Returns
grayim: float32 numpy array sized H x W with values in range [0, 1].
"""
grayim = cv2.imread(impath, 0)
if grayim is None:
raise Exception('Error reading image %s' % impath)
# Image is resized via opencv.
interp = cv2.INTER_AREA
grayim = cv2.resize(grayim, (img_size[1], img_size[0]), interpolation=interp)
grayim = (grayim.astype('float32') / 255.)
return grayim
def next_frame(self):
""" Return the next frame, and increment internal counter.
Returns
image: Next H x W image.
status: True or False depending whether image was loaded.
"""
if self.i == self.maxlen:
return (None, False)
if self.camera:
ret, input_image = self.cap.read()
if ret is False:
print('VideoStreamer: Cannot get image from camera (maybe bad --camid?)')
return (None, False)
if self.video_file:
self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.listing[self.i])
input_image = cv2.resize(input_image, (self.sizer[1], self.sizer[0]),
interpolation=cv2.INTER_AREA)
input_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2GRAY)
input_image = input_image.astype('float')/255.0
else:
image_file = self.listing[self.i]
input_image = self.read_image(image_file, self.sizer)
# Increment internal counter.
self.i = self.i + 1
input_image = input_image.astype('float32')
return (input_image, True)