FilesExpand file tree

 master

/

process.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

132 lines (111 loc) · 5.68 KB

 master

/

process.py

Copy path

Top

File metadata and controls

• Code
• Blame

132 lines (111 loc) · 5.68 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

# coding=utf-8

# Copyright 2019 The Google Research Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Forward processing of raw data to sRGB images.

Unprocessing Images for Learned Raw Denoising

http://timothybrooks.com/tech/unprocessing

"""

import numpy as np

import torch

import torch.nn as nn

import torch.distributions as tdist

def apply_gains(bayer_images, red_gains, blue_gains):

"""Applies white balance gains to a batch of Bayer images."""

red_gains = red_gains.squeeze()

blue_gains= blue_gains.squeeze()

bayer_images = bayer_images.permute(0, 2, 3, 1) # Permute the image tensor to BxHxWxC format from BxCxHxW format

green_gains = torch.ones_like(red_gains)

gains = torch.stack([red_gains, green_gains, green_gains, blue_gains], dim=-1)

#gains = gains[:, None, None, :]

gains = gains.reshape(1, 1, 1, 4)

outs = bayer_images * gains

outs = outs.permute(0, 3, 1, 2) # Re-Permute the tensor back to BxCxHxW format

return outs

def demosaic(bayer_images):

def SpaceToDepth_fact2(x):

# From here - https://discuss.pytorch.org/t/is-there-any-layer-like-tensorflows-space-to-depth-function/3487/14

bs = 2

N, C, H, W = x.size()

x = x.view(N, C, H // bs, bs, W // bs, bs) # (N, C, H//bs, bs, W//bs, bs)

x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # (N, bs, bs, C, H//bs, W//bs)

x = x.view(N, C * (bs ** 2), H // bs, W // bs) # (N, C*bs^2, H//bs, W//bs)

return x

def DepthToSpace_fact2(x):

# From here - https://discuss.pytorch.org/t/is-there-any-layer-like-tensorflows-space-to-depth-function/3487/14

bs = 2

N, C, H, W = x.size()

x = x.view(N, bs, bs, C // (bs ** 2), H, W) # (N, bs, bs, C//bs^2, H, W)

x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # (N, C//bs^2, H, bs, W, bs)

x = x.view(N, C // (bs ** 2), H * bs, W * bs) # (N, C//bs^2, H * bs, W * bs)

return x

"""Bilinearly demosaics a batch of RGGB Bayer images."""

bayer_images = bayer_images.permute(0, 2, 3, 1) # Permute the image tensor to BxHxWxC format from BxCxHxW format

shape = bayer_images.size()

shape = [shape[1] * 2, shape[2] * 2]

red = bayer_images[Ellipsis, 0:1]

upsamplebyX = nn.Upsample(size=shape, mode='bilinear', align_corners=False)

red = upsamplebyX(red.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)

green_red = bayer_images[Ellipsis, 1:2]

green_red = torch.flip(green_red, dims=[1]) # Flip left-right

green_red = upsamplebyX(green_red.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)

green_red = torch.flip(green_red, dims=[1]) # Flip left-right

green_red = SpaceToDepth_fact2(green_red.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)

green_blue = bayer_images[Ellipsis, 2:3]

green_blue = torch.flip(green_blue, dims=[0]) # Flip up-down

green_blue = upsamplebyX(green_blue.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)

green_blue = torch.flip(green_blue, dims=[0]) # Flip up-down

green_blue = SpaceToDepth_fact2(green_blue.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)

green_at_red = (green_red[Ellipsis, 0] + green_blue[Ellipsis, 0]) / 2

green_at_green_red = green_red[Ellipsis, 1]

green_at_green_blue = green_blue[Ellipsis, 2]

green_at_blue = (green_red[Ellipsis, 3] + green_blue[Ellipsis, 3]) / 2

green_planes = [

green_at_red, green_at_green_red, green_at_green_blue, green_at_blue

]

green = DepthToSpace_fact2(torch.stack(green_planes, dim=-1).permute(0, 3, 1, 2)).permute(0, 2, 3, 1)

blue = bayer_images[Ellipsis, 3:4]

blue = torch.flip(torch.flip(blue, dims=[1]), dims=[0])

blue = upsamplebyX(blue.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)

blue = torch.flip(torch.flip(blue, dims=[1]), dims=[0])

rgb_images = torch.cat([red, green, blue], dim=-1)

rgb_images = rgb_images.permute(0, 3, 1, 2) # Re-Permute the tensor back to BxCxHxW format

return rgb_images

def apply_ccms(images, ccms):

"""Applies color correction matrices."""

images = images.permute(0, 2, 3, 1) # Permute the image tensor to BxHxWxC format from BxCxHxW format

images = images[:, :, :, None, :]

#ccms = ccms[:, None, None, :, :]

outs = torch.sum(images * ccms, dim=-1)

outs = outs.permute(0, 3, 1, 2) # Re-Permute the tensor back to BxCxHxW format

return outs

def gamma_compression(images, gamma=2.2):

"""Converts from linear to gamma space."""

# Clamps to prevent numerical instability of gradients near zero.

images = images.permute(0, 2, 3, 1) # Permute the image tensor to BxHxWxC format from BxCxHxW format

outs = torch.clamp(images, min=1e-8) ** (1.0 / gamma)

outs = outs.permute(0, 3, 1, 2) # Re-Permute the tensor back to BxCxHxW format

return outs

def process(bayer_images, red_gains, blue_gains, cam2rgbs):

"""Processes a batch of Bayer RGGB images into sRGB images."""

# White balance.

bayer_images = apply_gains(bayer_images, red_gains, blue_gains)

# Demosaic.

bayer_images = torch.clamp(bayer_images, min=0.0, max=1.0)

images = demosaic(bayer_images)

# Color correction.

images = apply_ccms(images, cam2rgbs)

# Gamma compression.

images = torch.clamp(images, min=0.0, max=1.0)

images = gamma_compression(images)

return images