You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
119 lines
4.1 KiB
Python
119 lines
4.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
from scipy import signal
|
|
import torch
|
|
import os
|
|
|
|
|
|
PIL_CONVERT_COLOR = {
|
|
'RGB': lambda pil_image: pil_image.convert("RGB") if pil_image.mode != 'RGB' else pil_image,
|
|
'full_YCbCr': lambda pil_image: pil_image.convert("YCbCr") if pil_image.mode != 'YCbCr' else pil_image,
|
|
'full_Y': lambda pil_image: pil_image.convert("YCbCr").getchannel(0) if pil_image.mode != 'YCbCr' else pil_image.getchannel(0),
|
|
'sdtv_Y': lambda pil_image: _rgb2ycbcr(np.array(pil_image))[:,:,0] if pil_image.mode == 'RGB' else NotImplementedError(f"{pil_image.mode} to Y"),
|
|
'sdtv2_Y': lambda pil_image: rgb2y(np.array(pil_image)) if pil_image.mode == 'RGB' else NotImplementedError(f"{pil_image.mode} to Y"),
|
|
'L': lambda pil_image: pil_image.convert("L") if pil_image.mode != 'L' else pil_image,
|
|
}
|
|
|
|
def _rgb2ycbcr(img, maxVal=255):
|
|
O = np.array([[16],
|
|
[128],
|
|
[128]])
|
|
T = np.array([[0.256788235294118, 0.504129411764706, 0.097905882352941],
|
|
[-0.148223529411765, -0.290992156862745, 0.439215686274510],
|
|
[0.439215686274510, -0.367788235294118, -0.071427450980392]])
|
|
|
|
if maxVal == 1:
|
|
O = O / 255.0
|
|
|
|
t = np.reshape(img, (img.shape[0] * img.shape[1], img.shape[2]))
|
|
t = np.dot(t, np.transpose(T))
|
|
t[:, 0] += O[0]
|
|
t[:, 1] += O[1]
|
|
t[:, 2] += O[2]
|
|
ycbcr = np.reshape(t, [img.shape[0], img.shape[1], img.shape[2]])
|
|
|
|
return ycbcr
|
|
|
|
|
|
def rgb2y(im):
|
|
"""
|
|
this impl:
|
|
0.301, 0.586, 0.113 = 77/256, 150/256, 29/256
|
|
-0.172, -0.340, 0.512 = -44/256, -87/256, 131/256
|
|
0.512, -0.430, -0.082 = 131/256, -110/256, -21/256
|
|
|
|
ycbcr 601 sdtv spec[1]:
|
|
0.299, 0.587, 0.114
|
|
-0.172, -0.339, 0.511
|
|
0.511, -0.428, -0.083
|
|
|
|
ycbcr 601 sdtv spec[2]:
|
|
0.299, 0.587, 0.114
|
|
-0.169, -0.331, 0.5
|
|
0.5, -0.419, -0.081
|
|
|
|
[1] Video Demystified A Handbook for the Digital Engineer 4th ed - keith Jack, Chapter 3
|
|
[2] Color Space Conversions Adrian Ford, Alan Roberts
|
|
"""
|
|
im = im.astype(np.float32)
|
|
R, G, B = im[:,:,0], im[:,:,1], im[:,:,2]
|
|
Y = 77/256*R + 150/256*G + 29/256*B
|
|
# [1] Note that 8-bit YCbCr and R'G'B' data should be saturated a
|
|
return Y.clip(0,255).astype(np.uint8)
|
|
|
|
def rgb2yuv(im):
|
|
"""
|
|
this impl:
|
|
0.301, 0.586, 0.113 = 77/256, 150/256, 29/256
|
|
-0.172, -0.340, 0.512 = -44/256, -87/256, 131/256
|
|
0.512, -0.430, -0.082 = 131/256, -110/256, -21/256
|
|
|
|
ycbcr 601 sdtv spec[1]:
|
|
0.299, 0.587, 0.114
|
|
-0.172, -0.339, 0.511
|
|
0.511, -0.428, -0.083
|
|
|
|
ycbcr 601 sdtv spec[2]:
|
|
0.299, 0.587, 0.114
|
|
-0.169, -0.331, 0.5
|
|
0.5, -0.419, -0.081
|
|
|
|
[1] Video Demystified A Handbook for the Digital Engineer 4th ed - keith Jack, Chapter 3
|
|
[2] Color Space Conversions Adrian Ford, Alan Roberts
|
|
"""
|
|
im = im.astype(np.float32)
|
|
R, G, B = im[:,:,0], im[:,:,1], im[:,:,2]
|
|
Y = 77/256*R + 150/256*G + 29/256*B
|
|
U = -44/256*R - 87/256*G + 131/256*B
|
|
V = 131/256*R - 110/256*G - 21/256*B
|
|
Y, U, V = Y, U + 128, V + 128
|
|
# [1] Note that 8-bit YCbCr and R'G'B' data should be saturated at the 0 and 255 levels to avoid underflow and overflow
|
|
return np.stack([Y,U,V], axis=-1).clip(0,255).astype(np.uint8)
|
|
|
|
def yuv2rgb(im):
|
|
"""
|
|
this impl:
|
|
1, 0, 1.406 = 1, 0, 360/256
|
|
1, -0.344, -0.719 = 1, -88/256, -184/256
|
|
1, 1.777, 0 = 1, 455/256, 0
|
|
|
|
ycbcr 601 sdtv spec[1]:
|
|
1, 0, 1.371
|
|
1, -0.336, -0.698
|
|
1, 1.732, 0
|
|
|
|
ycbcr 601 sdtv spec[2]:
|
|
1, 0, 1.403
|
|
1, -0.344, -0.714
|
|
1, 1.773, 0
|
|
|
|
[1] Video Demystified A Handbook for the Digital Engineer 4th ed - keith Jack, Chapter 3
|
|
[2] Color Space Conversions Adrian Ford, Alan Roberts
|
|
"""
|
|
im = im.astype(np.float32)
|
|
Y, Ud, Vd = im[:,:,0], im[:,:,1]-128, im[:,:,2]-128
|
|
R = Y + 360/256*Vd
|
|
G = Y - 88/256*Ud - 184/256*Vd
|
|
B = Y + 455/256*Ud
|
|
# [1] Note that 8-bit YCbCr and R'G'B' data should be saturated at the 0 and 255 levels to avoid underflow and overflow
|
|
return np.stack([R, G, B], axis=-1).clip(0,255).astype(np.uint8) |