def test_read_without_init_start(self):
reader = TensorStreamConverter(self.path)
time.sleep(1.0)
with self.assertRaises(RuntimeError):
tensor, index = reader.read(return_index=True)
reader.stop()
def test_initialize_wrong_path(self):
reader = TensorStreamConverter("wrong.h264",
max_consumers=5,
cuda_device=0,
buffer_size=10)
with self.assertRaises(RuntimeError):
reader.initialize(repeat_number=5)
def test_initialize_correct_path(self):
reader = TensorStreamConverter(self.path,
max_consumers=5,
cuda_device=0,
buffer_size=10)
reader.initialize()
self.assertEqual(reader.frame_size, (1920, 1080))
self.assertEqual(reader.fps, 25)
def test_return_index(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.start()
time.sleep(1.0)
tensor, index = reader.read(return_index=True)
self.assertTrue(index > 0 and index < 100)
reader.stop()
def test_read_after_stop(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.start()
time.sleep(1.0)
reader.stop()
with self.assertRaises(RuntimeError):
tensor = reader.read()
def test_normalization(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.start()
time.sleep(1.0)
tensor = reader.read(normalization=True)
value = tensor[0][0][0].item()
self.assertEqual(type(value), float)
reader.stop()
def test_start_read_close(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.start()
time.sleep(1.0)
tensor = reader.read()
self.assertEqual(tensor.shape[0], 1080)
self.assertEqual(tensor.shape[1], 1920)
self.assertEqual(tensor.shape[2], 3)
reader.stop()
def test_constructor(self):
max_consumers = 5
cuda_device = 0
buffer_size = 10
reader = TensorStreamConverter(self.path,
max_consumers=max_consumers,
cuda_device=cuda_device,
buffer_size=buffer_size)
self.assertEqual(reader.max_consumers, max_consumers)
self.assertEqual(reader.cuda_device, cuda_device)
self.assertEqual(reader.buffer_size, buffer_size)
self.assertEqual(reader.stream_url, self.path)
def test_multiple_init(self):
reader = TensorStreamConverter(self.path)
number_close_init = 10
while number_close_init > 0:
reader.initialize()
reader.stop()
number_close_init -= 1
def test_constructor_default(self):
reader = TensorStreamConverter(self.path)
self.assertEqual(reader.max_consumers, 5)
self.assertEqual(reader.cuda_device, torch.cuda.current_device())
self.assertEqual(reader.buffer_size, 5)
self.assertEqual(reader.stream_url, self.path)
def test_close_start(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.stop()
#won't work but at least no crush
reader.start()
def test_start_close(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.start()
time.sleep(1.0)
reader.stop()
def test_logs_enabling(self):
reader = TensorStreamConverter(self.path)
reader.enable_logs(LogsLevel.LOW, LogsType.CONSOLE)
reader.enable_nvtx()
def test_stop_without_init(self):
reader = TensorStreamConverter(self.path)
reader.stop()
def test_check_dump_size(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.start()
time.sleep(1.0)
expected_width = 1920
expected_height = 1080
expected_channels = 3
tensor, index = reader.read(return_index=True)
self.assertEqual(tensor.shape[0], expected_height)
self.assertEqual(tensor.shape[1], expected_width)
self.assertEqual(tensor.shape[2], expected_channels)
# need to find dumped file and compare expected and real sizes
reader.dump(tensor)
dump_size = os.stat('default.yuv')
os.remove("default.yuv")
self.assertEqual(dump_size.st_size,
expected_width * expected_height * expected_channels)
reader.stop()
return model
def tensor_to_image(tensor):
image = tensor[0].to(torch.uint8)
image = image.permute(1, 2, 0)
image = image.cpu().numpy()
return image
if __name__ == "__main__":
args = parse_arguments()
style_model = load_model(args.model, device='cuda')
reader = TensorStreamConverter(args.input, repeat_number=20)
reader.initialize()
print(f"Input video frame size: {reader.frame_size}, fps: {reader.fps}")
width = args.width if args.width else reader.frame_size[0]
height = args.height if args.height else reader.frame_size[1]
print(f"Model input image width: {width}, height: {height}")
writer = FFmpegVideoWriter(
args.output,
out_size=(width * 2 if args.concat_orig else width, height),
out_fps=reader.fps,
bitrate=args.bitrate,
codec=args.codec,
preset=args.preset)
def test_frame_number(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.start()
time.sleep(1.0)
frame_num = i = 10
while i > 0:
tensor = reader.read()
reader.dump(tensor)
i -= 1
dump_size = os.stat('default.yuv')
print(f"SIZE {dump_size}")
os.remove("default.yuv")
expected_width = 1920
expected_height = 1080
expected_channels = 3
expected_size = expected_width * expected_height * expected_channels * frame_num
self.assertEqual(dump_size.st_size, expected_size)
reader.stop()
def test_dump_name(self):
reader = TensorStreamConverter(self.path)
reader.initialize()
reader.start()
time.sleep(1.0)
tensor = reader.read()
# need to find dumped file and compare expected and real sizes
reader.dump(tensor, name="dump")
self.assertTrue(os.path.isfile("dump.yuv"))
os.remove("dump.yuv")
reader.stop()
"--verbose",
default="LOW",
choices=["LOW", "MEDIUM", "HIGH"],
help="Set output level from library (default: LOW)")
parser.add_argument("-n",
"--number",
help="Number of frame to parse (default: unlimited)",
type=int,
default=0)
return parser.parse_args()
if __name__ == '__main__':
args = parse_arguments()
reader = TensorStreamConverter(args.input, repeat_number=20)
reader.enable_logs(LogsLevel[args.verbose], LogsType.CONSOLE)
reader.initialize()
reader.start()
if args.output:
if os.path.exists(args.output):
os.remove(args.output)
tensor = None
try:
while True:
tensor, index = reader.read(pixel_format=FourCC[args.fourcc],
return_index=True,
width=args.width,
reader.dump(tensor, args.output2, **parameters)
if index % int(reader.fps) == 0:
print("consumer2 frame index", index)
reader.stop()
time.sleep(1.0) # prevent simultaneous print
print("consumer2 shape:", tensor.shape)
print("consumer2 dtype:", tensor.dtype)
print("consumer2 last frame index:", index)
if __name__ == "__main__":
args = parse_arguments()
reader1 = TensorStreamConverter(args.input1, cuda_device=args.cuda_device1)
reader1.enable_logs(LogsLevel[args.verbose1], LogsType.CONSOLE)
reader1.initialize(repeat_number=20)
reader2 = TensorStreamConverter(args.input2, cuda_device=args.cuda_device2)
reader2.enable_logs(LogsLevel[args.verbose2], LogsType.CONSOLE)
reader2.initialize(repeat_number=20)
reader1.start()
reader2.start()
thread1 = Thread(target=consumer1, args=(reader1, args.number1))
thread2 = Thread(target=consumer2, args=(reader2, args.number2))
thread1.start()
thread2.start()