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
|
import torch
from torchaudio._internal.module_utils import is_module_available
if is_module_available("PIL"):
from PIL import Image
def save_image(path, data, mode=None):
"""Save image.
The input image is expected to be CHW order
"""
if torch.is_tensor(data):
data = data.numpy()
if mode == "L" and data.ndim == 3:
assert data.shape[0] == 1
data = data[0]
if data.ndim == 3:
data = data.transpose(1, 2, 0)
Image.fromarray(data, mode=mode).save(path)
def get_image(width, height, grayscale=False):
"""Generate image Tensor, returns CHW"""
channels = 1 if grayscale else 3
numel = width * height * channels
img = torch.arange(numel, dtype=torch.int64) % 256
img = img.reshape(channels, height, width).to(torch.uint8)
return img
def rgb_to_yuv_ccir(img):
"""rgb to yuv conversion ported from ffmpeg
The input image is expected to be (..., channel, height, width).
"""
assert img.dtype == torch.uint8
img = img.to(torch.float32)
r, g, b = torch.split(img, 1, dim=-3)
# https://github.com/FFmpeg/FFmpeg/blob/870bfe16a12bf09dca3a4ae27ef6f81a2de80c40/libavutil/colorspace.h#L98
y = 263 * r + 516 * g + 100 * b + 512 + 16384
y /= 1024
# https://github.com/FFmpeg/FFmpeg/blob/870bfe16a12bf09dca3a4ae27ef6f81a2de80c40/libavutil/colorspace.h#L102
# shift == 0
u = -152 * r - 298 * g + 450 * b + 512 - 1
u /= 1024
u += 128
# https://github.com/FFmpeg/FFmpeg/blob/870bfe16a12bf09dca3a4ae27ef6f81a2de80c40/libavutil/colorspace.h#L106
# shift == 0
v = 450 * r - 377 * g - 73 * b + 512 - 1
v /= 1024
v += 128
return torch.cat([y, u, v], -3).to(torch.uint8)
def rgb_to_gray(img):
"""rgb to gray conversion
The input image is expected to be (..., channel, height, width).
"""
assert img.dtype == torch.uint8
img = img.to(torch.float32)
r, g, b = torch.split(img, 1, dim=-3)
gray = 0.299 * r + 0.587 * g + 0.114 * b
return gray.to(torch.uint8)
|
