diff --git a/imutils/video/filevideostream.py b/imutils/video/filevideostream.py index f173eb5..a2bc6f4 100644 --- a/imutils/video/filevideostream.py +++ b/imutils/video/filevideostream.py @@ -6,33 +6,65 @@ # import the Queue class from Python 3 if sys.version_info >= (3, 0): - from queue import Queue + from queue import Queue, Full # otherwise, import the Queue class for Python 2.7 else: from Queue import Queue +class EndOfStreamException(Exception): + pass + + class FileVideoStream: - def __init__(self, path, transform=None, queue_size=128): + + def __init__(self, path, transform=None, queue_size=128, skip_frames=True, queue_put_timeout=0.5): # initialize the file video stream along with the boolean # used to indicate if the thread should be stopped or not + self.queue_put_timeout = queue_put_timeout + self.skip_frames = skip_frames self.stream = cv2.VideoCapture(path) + self.total_frames = self.stream.get(cv2.CAP_PROP_FRAME_COUNT) + self.skipped_frames = 0 self.stopped = False self.transform = transform # initialize the queue used to store frames read from # the video file self.Q = Queue(maxsize=queue_size) - # intialize thread + # initialize thread self.thread = Thread(target=self.update, args=()) self.thread.daemon = True + def check_eos(self): + if self.stream.get(cv2.CAP_PROP_POS_FRAMES) == self.total_frames: + # If the number of captured frames is equal to the total number of frames, + # we stop + self.stop() + raise EndOfStreamException("End of stream") + def start(self): # start a thread to read frames from the file video stream self.thread.start() return self + def get_frame(self): + # grab the current frame + flag, frame = self.stream.read() + + while not flag and self.skip_frames: + self.check_eos() + # if self.skipped_frames == 0: + # print(f"Skipping frame(s)") + self.skipped_frames += 1 + flag, frame = self.stream.read() + if self.skipped_frames > 0: + # print(f"Resuming video...") + self.skipped_frames = 0 + + return flag, frame + def update(self): # keep looping infinitely while True: @@ -40,24 +72,22 @@ def update(self): # thread if self.stopped: break - - # otherwise, ensure the queue has room in it - if not self.Q.full(): + try: # read the next frame from the file - (grabbed, frame) = self.stream.read() - + grabbed, frame = self.get_frame() # if the `grabbed` boolean is `False`, then we have # reached the end of the video file if not grabbed: self.stopped = True - + raise EndOfStreamException() + # if there are transforms to be done, might as well # do them on producer thread before handing back to - # consumer thread. ie. Usually the producer is so far + # consumer thread. i.e. Usually the producer is so far # ahead of consumer that we have time to spare. # # Python is not parallel but the transform operations - # are usually OpenCV native so release the GIL. + # are typically OpenCV native so release the GIL. # # Really just trying to avoid spinning up additional # native threads and overheads of additional @@ -65,13 +95,26 @@ def update(self): # idle grabbing frames. if self.transform: frame = self.transform(frame) + while True: + try: + # try to put it into the queue + self.Q.put(frame, True, self.queue_put_timeout) + break + except Full: + # after timeout seconds check if we are still not stopped + if self.stopped: + raise EndOfStreamException() + except EndOfStreamException: + # if we got stopped or reached eos + self.stopped = True + break + self.stream.release() - # add the frame to the queue - self.Q.put(frame) - else: - time.sleep(0.1) # Rest for 10ms, we have a full queue + def dim(self): + width = int(self.stream.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(self.stream.get(cv2.CAP_PROP_FRAME_HEIGHT)) - self.stream.release() + return [width, height] def read(self): # return next frame in the queue