def Pred(): cv2.namedWindow("cnn") a = win32gui.FindWindow(None, "cnn") win32gui.SetWindowLong(a, win32con.GWL_STYLE, win32con.WS_POPUP) x = 0 y = 0 width = 1920 height = 1080 while True: start_time = time.perf_counter() img = fB.get(timeout=2) frame = img frame = torch.from_numpy(frame).to(dev).reshape( [1, 1080, 1920, 3]).permute(0, 3, 1, 2).float() #nchw frameYUV = kornia.rgb_to_yuv(frame) frameY = frameYUV[:, 0, :, :].reshape([1, 1, 1080, 1920]).half() predY = model(frameY / 255) * 255 f0 = predY.float()[0, 0, :, :] f1 = frameYUV[0, 1, :, :] f2 = frameYUV[0, 2, :, :] #1080,1920 pred = torch.stack([f0, f1, f2]).float() # 3,1080,1920 pred = kornia.yuv_to_rgb(pred).clip(0, 255).permute(1, 2, 0)[:, :, [2, 1, 0]] pred = torch.tensor(pred, dtype=torch.uint8).to('cpu').numpy() cv2.imshow('cnn', pred) win32gui.SetWindowPos(a, win32con.HWND_TOPMOST, x, y, width, height, win32con.SWP_SHOWWINDOW) if cv2.waitKey(1) & 0xFF == ord('q'): break x1 = time.perf_counter() - start_time start_timeD = time.perf_counter()
def test_yuv_to_rbg_shape_bad(self, bad_input_shapes): with pytest.raises(ValueError): out = kornia.yuv_to_rgb(torch.ones(*bad_input_shapes))
def test_yuv_to_rbg_type(self): with pytest.raises(TypeError): out = kornia.yuv_to_rgb(1)
def test_yuv_to_rgb_batch_shape(self, device): batch_size, channels, height, width = 2, 3, 4, 5 img = torch.ones(batch_size, channels, height, width).to(device) assert kornia.yuv_to_rgb(img).shape == \ (batch_size, channels, height, width)
def test_yuv_to_rgb_shape(self, device): channels, height, width = 3, 4, 5 img = torch.ones(channels, height, width).to(device) assert kornia.yuv_to_rgb(img).shape == (channels, height, width)