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()
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],
'resize_type': ResizeType[args.resize_type]
}
tensor, index = reader.read(**parameters, return_index=True)
if args.number:
if index > args.number:
break
if args.output:
reader.dump(tensor, args.output, **parameters)
except RuntimeError as e:
print(f"Bad things happened: {e}")
finally:
print("Frame size: ", reader.frame_size)
print("FPS: ", reader.fps)
if tensor is not None:
print("Tensor shape:", tensor.shape)
print("Tensor dtype:", tensor.dtype)
print("Tensor device:", tensor.device)
reader.stop()
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,
height=args.height)
if args.number:
if index > args.number:
break
if args.output:
reader.dump(tensor, args.output)
except RuntimeError as e:
print(f"Bad things happened: {e}")
finally:
print("Frame size: ", reader.frame_size)
print("FPS: ", reader.fps)
if tensor is not None:
print("Tensor shape:", tensor.shape)
print("Tensor dtype:", tensor.dtype)
print("Tensor device:", tensor.device)
reader.stop()
