TensorRT-Demo/utils/background.py

66 lines
2.2 KiB
Python
Raw Normal View History

2023-03-06 20:44:29 +08:00
"""background.py
This code implements the Background class for the TensorRT MODNet
demo. The Background class could generate background images from
either a still image, a video file or nothing (pure black bg).
"""
import numpy as np
import cv2
class Background():
"""Backgrounf class which supports one of the following sources:
1. Image (jpg, png, etc.) file, repeating indefinitely
2. Video file, looping forever
3. None -> black background
# Arguments
src: if not spcified, use black background; else, src should be
a filename of an image (jpg/png) or video (mp4/ts)
width & height: width & height of the output background image
"""
def __init__(self, src, width, height, demo_mode=False):
self.src = src
self.width = width
self.height = height
self.demo_mode = demo_mode
if not src: # empty source: black background
self.is_video = False
self.bg_frame = np.zeros((height, width, 3), dtype=np.uint8)
elif not isinstance(src, str):
raise ValueError('bad src')
elif src.endswith('.jpg') or src.endswith('.png'):
self.is_video = False
self.bg_frame = cv2.resize(cv2.imread(src), (width, height))
assert self.bg_frame is not None and self.bg_frame.ndim == 3
elif src.endswith('.mp4') or src.endswith('.ts'):
self.is_video = True
self.cap = cv2.VideoCapture(src)
assert self.cap.isOpened()
else:
raise ValueError('unknown src')
def read(self):
"""Read a frame from the Background object."""
if self.is_video:
_, frame = self.cap.read()
if frame is None:
# assume end of video file has been reached, so loop around
self.cap.release()
self.cap = cv2.VideoCapture(self.src)
_, frame = self.cap.read()
return cv2.resize(frame, (self.width, self.height))
else:
return self.bg_frame.copy()
def __del__(self):
if self.is_video:
try:
self.cap.release()
except:
pass