forked from Breakthrough/PySceneDetect
-
Notifications
You must be signed in to change notification settings - Fork 0
/
bbox_video_stream.py
118 lines (102 loc) · 4.39 KB
/
bbox_video_stream.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
111
112
113
114
115
116
117
118
# -*- coding: utf-8 -*-
#
""" ``scenedetect.bbox_video_stream`` Module.
This module contains the BboxVideoStream class that sets up a video stream
and allows the user to select a recatangle / bounding box (coordinates
stored in a list of int in format: [ul-x, ul-y, lr-x, lr-y]) by mouse
clicking upper left and lower right coordinates of the desired bounding box
within the video frame.
The bbox ccordinates are saved to a .bbox text_file and can then be used to
crop a cv2 frame.
"""
# Standard Library Imports
import logging
import os
import sys
# Third-Party Library Imports
import click
import cv2
class BboxVideoStream(object):
""" BboxVideoStream: base class that sets up a video stream and allows the
user to select a recatangle / bounding box. """
def __init__(self, video_file):
""" BboxVideoStream Constructor Method (__init__).
Arguments:
video_file (str): Absolute path to video file.
Raises:
VideoOpenFailure: Video could not be opened.
"""
self.cap = cv2.VideoCapture(video_file)
self.output_file = os.path.splitext(video_file)[0] + ".bbox"
self.basename = os.path.basename(video_file)
self.font = cv2.FONT_HERSHEY_SIMPLEX
self.bbox = []
if (self.cap.isOpened()== False):
raise VideoOpenFailure("Error opening the video file.")
return
else:
self.frame = self.cap.read()[1]
self.clean_frame = None
self.height = self.frame.shape[0]
def mouse_callback(self, event, x, y, flags, params):
""" Callback method that sets bbox (rect) coordinates on an image. """
if event == cv2.EVENT_LBUTTONDOWN:
if len(self.bbox) == 4:
del self.bbox[:]
self.frame = self.clean_frame
cv2.imshow(self.basename, self.frame)
self.bbox.extend([x, y])
if len(self.bbox) == 2:
cv2.circle(self.frame, tuple(self.bbox), 2, (0,255,0), 2)
cv2.imshow(self.basename, self.frame)
elif len(self.bbox) == 4:
self._add_bbox_overlay()
cv2.imshow(self.basename, self.frame)
self._check_and_fix_bbox()
self._save_bbox()
def _add_bbox_overlay(self):
""" Adds a rectangle overlay with the bbox coordinates to the frame. """
cv2.rectangle(self.frame,
tuple(self.bbox[:2]), tuple(self.bbox[2:]),
(0,255,0), 2)
def _add_text_overlay(self, text, y_pos):
""" Adds a text overlay to the cv2 frame. """
cv2.putText(self.frame, text, (30, y_pos), self.font, 1, (0, 255, 0), 2,
cv2.LINE_AA)
def _check_and_fix_bbox(self):
""" Makes sure that bbox has this format: [ul-x, ul-y, lr-x, lr-y]."""
if self.bbox[0] > self.bbox[2]:
self.bbox[0], self.bbox[2] = self.bbox[2], self.bbox[0]
if self.bbox[1] > self.bbox[3]:
self.bbox[1], self.bbox[3] = self.bbox[3], self.bbox[1]
def _save_bbox(self):
""" Saves bbox to textfile. """
with open(self.output_file, "w") as text_file:
text_file.write(str(self.bbox))
logging.info('Bounding Box: %s saved to %s', str(self.bbox),
self.output_file)
def play(self):
""" Plays video and lets user define the bbox after pressing any key."""
while(self.cap.isOpened()):
status, self.frame = self.cap.read()
info_text = "Press any key to pause video. Press q to exit."
self._add_text_overlay(info_text, int(self.height/2))
if status:
if len(self.bbox) == 4:
self.clean_frame = self.frame.copy()
self._add_bbox_overlay()
cv2.imshow(self.basename, self.frame)
key = cv2.waitKey(1)
if key == ord("q"):
break
elif key > -1:
info_text = ("Set Bounding Box by clicking Upper Left & Lower"
"Right.")
self._add_text_overlay(info_text, int(self.height/2-50))
cv2.imshow(self.basename, self.frame)
cv2.setMouseCallback(self.basename, self.mouse_callback)
cv2.waitKey(0)
else:
break
self.cap.release()
cv2.destroyAllWindows()