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_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_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_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_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_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()
args = parse_arguments()
reader = TensorStreamConverter(
args.input,
max_consumers=5,
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.BGR24,
return_index=True)
if index % int(reader.fps) == 0:
print("consumer2 frame index", index)
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()
reader = TensorStreamConverter(args.input, repeat_number=20)
reader.initialize()
reader.start()
thread1 = Thread(target=consumer1, args=(reader, args.number))
thread2 = Thread(target=consumer2, args=(reader, args.number))
thread1.start()
thread2.start()
thread1.join()
thread2.join()
reader.stop()
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()