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_read_without_init(self):
reader = TensorStreamConverter(self.path)
reader.start()
time.sleep(1.0)
with self.assertRaises(RuntimeError):
tensor, index = reader.read(return_index=True)
reader.stop()
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_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()
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()
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()
cuda_device=args.cuda_device,
buffer_size=args.buffer_size,
framerate_mode=FrameRate[args.framerate_mode])
#To log initialize stage, logs should be defined before initialize call
reader.enable_logs(LogsLevel[args.verbose],
LogsType[args.verbose_destination])
if args.nvtx:
reader.enable_nvtx()
reader.initialize(repeat_number=20)
if args.skip_analyze:
reader.skip_analyze()
reader.start()
if args.output:
if os.path.exists(args.output + ".yuv"):
os.remove(args.output + ".yuv")
print(f"Normalize {args.normalize}")
tensor = None
try:
while True:
parameters = {
'pixel_format': FourCC[args.fourcc],
'width': args.width,
'height': args.height,
'normalization': args.normalize,
'planes_pos': Planes[args.planes],
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()
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()
thread1.join()
thread2.join()
