File size: 406 Bytes
0145b71 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
import torch
def cv2ten(img, device):
img = (img[:, :, ::-1].transpose(2, 0, 1) / 255. - 0.5) / 0.5
img_ten = torch.from_numpy(img).float().unsqueeze(0).to(device)
return img_ten
def ten2cv(img_ten, bgr=True):
img = img_ten.squeeze(0).mul_(0.5).add_(0.5).mul_(255).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
if bgr:
img = img[:, :, ::-1]
return img
|