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()