자동 화이트 밸런싱에 대한 간단하고 가벼운 접근 방식입니다.

지난 몇 달 동안 저는 브라우저 내 차별화 가능한 사진 편집기를 구축했습니다. 목표는 각 이미지 처리 기능을 기계 학습 모델로 쉽게 매개변수화하거나 학습된 변환으로 만드는 것입니다. 다운스트림 작업에 대한 손실 함수에서 각각을 사용할 수 있기를 원합니다. 오늘 저는 자동 화이트 밸런스 보정 접근 방식(및 기타 전역 이미지 변환)을 공유하고 싶었습니다.

나는 주로 회색 세계 및 YUV 좌표 기반 변환보다 더 나은 성능을 발휘하는 효율적인 매개변수(4k 매개변수!)와 빠른 접근 방식을 찾고 있기 때문에 최신 결과에 대해 어려운 주장을 하지 않을 것입니다. 그러나 예비 테스트에서는 U-NET 기반 접근 방식에 비해 꽤 잘 작동합니다. 중요한 것은 하이퍼파라미터 조정 없이 GTX 1070에서 몇 시간 만에 훈련할 수 있다는 것입니다. 또한 경량 장치의 브라우저에서도 성능을 발휘하는데 이것이 저의 주요 사용 사례입니다.

이제 코드로 이동하겠습니다.

class WhiteBalance(nn.Module):
    def __init__(self, hidden_dim=10):
        super(WhiteBalance, self).__init__()
        self.parameter_network = nn.Sequential(
            nn.Conv2d(3, 7, kernel_size=(3, 3), padding=1),
            nn.MaxPool2d(kernel_size=2),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Conv2d(7, 14, kernel_size=(3, 3), padding=1),
            nn.MaxPool2d(kernel_size=2),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Conv2d(14, 3, kernel_size=(3, 3), padding=1),
        )
        self.feature_dim = 192
        self.output_dim = 96 # Enough Parameters for a 16 dim inner neural network
        self.weight_size = self.output_dim//2
        self.parameter_network.apply(weights_init)

        self.linear = nn.Sequential(
            nn.Linear(self.feature_dim, hidden_dim),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Linear(hidden_dim, self.output_dim),
        )

    def inner_network(self, image, params):

        batch, c, h, w = image.size()
        # split parameters for each layer of the inner network
        w_1, w_2 = params[:, :self.weight_size], params[:, self.weight_size:]
        hidden_dim_inner = self.weight_size//c
        # Reshape matrix into hidden_dim x channels
        w_1, w_2 = w_1.view(-1, hidden_dim_inner, c), w_2.view(-1, c, hidden_dim_inner)
        # apply batch of neural networks to image pixelwise
        pixels = image.permute(0, 2, 3, 1).contiguous().view(batch, -1, c)
        output = torch.selu(pixels.bmm(w_1.permute(0, 2, 1)))
        output = output.bmm(w_2.permute(0, 2, 1)).view(-1, h, w, c).permute(0, 3, 1, 2)
        return output

    def forward(self, input_image, full_resolution_image):
        batch = input_image.size(0)
        # Output Parameters to a small inner neural network using downsampled input image or patch.
        params = self.linear(self.parameter_network(input_image).view(batch,  self.feature_dim))
        # Apply pixelwise neural network to full scale image
        output = self.inner_network(full_resolution_image, params)
        # Scale in [0, 1]
        return torch.clamp(output, 0, 1)

따라서 아키텍처에 대해 조금 알아보면 이 모델은 32x32 sRGB 이미지를 입력으로 사용하고 숨겨진 차원이 16이고 입력 및 출력 크기가 3인 작은 2계층 신경망에 대한 매개변수를 출력합니다. 이것이 의미하는 바는 다음과 같습니다. 모든 해상도의 이미지에 적용할 수 있는 픽셀 단위 변환을 출력합니다.

이것에는 많은 이점이 있습니다.

매개변수를 WebGL 셰이더에 쉽게 전달할 수 있습니다.

각 픽셀 색상을 점진적으로 공급하여 3D 조회 테이블(LUT)을 생성할 수 있으므로 실시간으로 비디오에 동일한 화이트 밸런싱을 적용할 수 있습니다.

우리가 배울 수 있는 비선형 픽셀 단위 변환 유형의 측면에서 매우 유연합니다.

교육이 진행되는 한 데이터 세트는 MAE 손실 또는 이와 유사한 흰색 균형 사진 쌍 전후입니다. 매개변수 네트워크에 대한 입력보다 더 높은 분해능으로 훈련함으로써 성능이 향상되는 것을 확인했습니다. 예를 들어 입력으로 32x32 이미지가 있을 수 있습니다. 그러나 224x224 버전에서 손실을 계산합니다. 정규화 효과가 있는 것 같습니다.

이런 식으로 생각하십시오 ...

input_image, larger_image, larger_target = batch
output = model(input_image, larger_image)
loss = F.l1_loss(output, target)

다음은 몇 가지 샘플입니다. (중복 무시, 약간의 입력 변형 테스트)
전에

후에


저에너지 장치에 대한 옵션을 찾고 있는 경우 속도를 높이고 매개변수를 더 줄일 수 있는 방법이 있습니다. 예를 들어 매개변수 네트워크의 출력에 대해 컨벌루션 채널별 평균 풀링을 수행할 수 있습니다(아마도 학습된 가중치 사용?). 그러면 선형 네트워크와 내부 신경망을 제거할 수 있습니다.

그러나 저는 사용 사례에 대해 전역 비선형 픽셀 단위 변환의 유연성을 선호합니다.

읽어주셔서 감사합니다. 이런 종류의 자료에 관심이 있으시면 제 .