File size: 3,495 Bytes
a220803 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import numpy as np
import cv2
import torch
def normalize_tensor(in_feat,eps=1e-10):
norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1,keepdim=True))
return in_feat/(norm_factor+eps)
def l2(p0, p1, range=255.):
return .5*np.mean((p0 / range - p1 / range)**2)
def dssim(p0, p1, range=255.):
from skimage.measure import compare_ssim
return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.
def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
image_numpy = image_tensor[0].cpu().float().numpy()
image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
return image_numpy.astype(imtype)
def tensor2np(tensor_obj):
# change dimension of a tensor object into a numpy array
return tensor_obj[0].cpu().float().numpy().transpose((1,2,0))
def np2tensor(np_obj):
# change dimenion of np array into tensor array
return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
def tensor2tensorlab(image_tensor,to_norm=True,mc_only=False):
# image tensor to lab tensor
from skimage import color
img = tensor2im(image_tensor)
img_lab = color.rgb2lab(img)
if(mc_only):
img_lab[:,:,0] = img_lab[:,:,0]-50
if(to_norm and not mc_only):
img_lab[:,:,0] = img_lab[:,:,0]-50
img_lab = img_lab/100.
return np2tensor(img_lab)
def read_frame_yuv2rgb(stream, width, height, iFrame, bit_depth, pix_fmt='420'):
if pix_fmt == '420':
multiplier = 1
uv_factor = 2
elif pix_fmt == '444':
multiplier = 2
uv_factor = 1
else:
print('Pixel format {} is not supported'.format(pix_fmt))
return
if bit_depth == 8:
datatype = np.uint8
stream.seek(iFrame*1.5*width*height*multiplier)
Y = np.fromfile(stream, dtype=datatype, count=width*height).reshape((height, width))
# read chroma samples and upsample since original is 4:2:0 sampling
U = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
reshape((height//uv_factor, width//uv_factor))
V = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
reshape((height//uv_factor, width//uv_factor))
else:
datatype = np.uint16
stream.seek(iFrame*3*width*height*multiplier)
Y = np.fromfile(stream, dtype=datatype, count=width*height).reshape((height, width))
U = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
reshape((height//uv_factor, width//uv_factor))
V = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
reshape((height//uv_factor, width//uv_factor))
if pix_fmt == '420':
yuv = np.empty((height*3//2, width), dtype=datatype)
yuv[0:height,:] = Y
yuv[height:height+height//4,:] = U.reshape(-1, width)
yuv[height+height//4:,:] = V.reshape(-1, width)
if bit_depth != 8:
yuv = (yuv/(2**bit_depth-1)*255).astype(np.uint8)
#convert to rgb
rgb = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB_I420)
else:
yvu = np.stack([Y,V,U],axis=2)
if bit_depth != 8:
yvu = (yvu/(2**bit_depth-1)*255).astype(np.uint8)
rgb = cv2.cvtColor(yvu, cv2.COLOR_YCrCb2RGB)
return rgb
|