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_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_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()
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],
'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)
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)
reader.start()
try:
while True:
tensor, index = reader.read(pixel_format=FourCC.RGB24,
return_index=True,
width=width,
height=height,
planes_pos=Planes.PLANAR,
normalization=True)
tensor = tensor.unsqueeze(0)
with torch.no_grad():
output = style_model(tensor)
output = torch.clamp(output, 0, 255)
style_frame = tensor_to_image(output)
if args.concat_orig:
orig_frame = tensor_to_image(tensor)
write_frame = np.concatenate((orig_frame, style_frame), axis=1)
else:
write_frame = style_frame
writer.write(write_frame)
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,
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(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)
reader.start()
try:
while True:
tensor, index = reader.read(pixel_format=FourCC.RGB24,
return_index=True,
width=width,
height=height)
tensor = tensor.permute(2, 0, 1)
tensor = tensor.unsqueeze(0)
tensor = tensor.to(torch.float32)
with torch.no_grad():
output = style_model(tensor)
output = torch.clamp(output, 0, 255)
style_frame = tensor_to_image(output)
if args.concat_orig:
orig_frame = tensor_to_image(tensor)
write_frame = np.concatenate((orig_frame, style_frame), axis=1)
else: