😊AI를 이용해서 어두운 이미지를 밝게 만들어보자

[gellston]

육아 휴직 중인데 심심하군요..

AI를 이용해서 어두운 이미지를 밝게 만드는 프로그램을 만들어보고 싶어졌습니다.

아래 이미지는 저조도 데이터셋 이미지인데 재미있어 보이는군요.

조금씩 만들어보면서 결과들을 업데이트 해보겠습니다.

[gellston]

가장 만만해 보이는 논문을 찾았습니다.

https://arxiv.org/pdf/2103.00860

요 녀석을 정독해 보면서 아래와 같은 순서로 구현을 해볼까 합니다.

1. pytorch를 이용해서 모델을 만듬
2. C++에서 돌릴 수 있게 inference 코드를 작성
3. C#으로 그것을 다시 랩핑 (마샬링방법과 C++/CLI 두가지 방법으로 래핑을 진행해서 장단점 정리)
4. C# WPF에 서비스로 밀어 넣어서 간단하게 이미지를 불러와 입력 이미지와 결과 이미지를 비교

근데.. 버릇이 참 무서운게 분명 휴가 중인데 눈이 로봇처럼 새벽이 되면 열리는군요..

[gellston]

너는 도대체 어떻게 동작 하는 거니?

Zero-DCE모델은 읽어 보니.. 입력 RGB 입력을 받으면 그 입력으로 딥러닝 모델을 통과해 추정한 곡선 커브용 파라미터와 딥러닝 모델의 입력으로 넣었던 x와 함께 곱하거나 더해서 최적의 밝기 곡선을 가진 이미지를 만드는것이 목적입니다.

논문 저자는 8번 반복해서 곱해서 최적의 밝기를 가진 이미지를 만들어냈습니다.

아.. 그리고 음.. 논문저자의 원문 repository도 참고했는데 저 커브 파라미터는 입력이미지가 500x500이라면 그 갯수만큼 따로 추정합니다. 이해가 안되서 한참을 봤는데… 픽셀마다 파라미터 weight가 독립이니 어두운 곳은 밝게 밝은 곳은 더 작은 가중치를 줘서 번지지 않게 할 수 있을 것 같다는 느낌적인 느낌이 드는군요. 만약 저 파라미터가 하나였다면 당연하겠지만 밝은곳과 어두운 곳을 따로 컨트롤 할 수 없겠지요 ㅠ

위에 설명한 것 처럼 모든 픽셀이 다른 가중치를 가지기 때문에 딥러닝이 조정한 파라미터 값들을 R,G,B 채널로 분리해서 시각화하여 heatmap으로 표현했을 때 어두운 곳은 많은 가중치, 밝은 곳은 적은 가중치를 곱하게 됩니다.

딥러닝 모델의 구조적 관점으로 관찰해보자

이건 뭐.. 전혀 복잡하지 않은 여러 컨볼루션을 적층하여 만든 단순한 모델이군요.

다만 아무리 단순하더라도 stride나 max pooling을 이용한 피쳐맵 축소가 없어서 메모리가 터져버릴 것 같은데.. 읽어보니 저 컨볼루션은 단순 컨볼루션이 아니라 google이 mobilenet이랑 xception모델에서 사용한 separable convolution을 이용했습니다.

separable convolution에 대한 설명을 적으면 너무 길어질 것 같아서 아래 링크로 대체하겠습니다.

블로그 운영하시는 분들이 글을 정말 잘 쓰시는군요 ….

Depthwise separable convolution 연산

그 외에 특별한 기술적인 부분은 없고 연산량을 줄이기 위해서 커브 파라미터 추정 시에 사용자가 지정한 scale값 만큼을 줄이는 작업을 하는군요.

코드를 뜯어보니.. 가로 세로 길이를 4배를 줄입니다. 면적으로 따지면 16분의 1로 줄어 들었다고 볼 수 있겠군요.

코딩하는 개발자가 바라보는 동작 원리

저는 딥러닝에 관심이 있는 사람이지 대학원을 다니거나 연구원은 아니기 때문에 수학을 정말 못합니다 -_-

오히려 C# 코딩이 좋은데 그럼 저 모델을 도대체 어떻게 이해해야 될 가 생각해봤습니다.

딱 맞는 정확한 설명은 아니지만 아래 그림처럼 직관적인 느낌이 그려지는군요..

Slog 태클과 의견공유… 댓글은 언제나 환영입니다.

[gellston]

구현한 저조도 개선 모델 Python코드

논문에 나온대로 구현하였고 학습을 걸어놓았습니다.
결과가 잘 나왔으면 좋겠군요..
이전 포스트에서 언급했듯이 이 모델을 학습시켜서 배포는 ONNX로 진행을 할 것입니다.
그 ONNX모델을 C++ ONNX라이브러리를 이용하여 로드해서 돌아가도록 만들 것이고.
그 이후에 C# 래핑 과정을 거쳐 UI에 시각적으로 표시할 예정입니다.
물론 누겟에 배포되어있는 유명한 ONNX C#라이브러리를 이용해서 추론하는것도 가능합니다.

Microsoft.ML.OnnxRuntime 1.23.2

아 그리고 제가 사용한 저조도 개선 데이터셋은 아래 링크에서 받으실 수 있습니다.

LOL Dataset

Separable Convolution Block

import torch
class DepthwiseSeparableConv(torch.nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DepthwiseSeparableConv, self).__init__()
        self.depth_conv = torch.nn.Conv2d(in_ch, in_ch, kernel_size=3, padding=1, groups=in_ch)
        self.point_conv = torch.nn.Conv2d(in_ch, out_ch, kernel_size=1)

    def forward(self, x):
        return self.point_conv(self.depth_conv(x))

ZeroDCE Model

import torch
import torch.nn as nn
import torch.nn.functional as F



from model.block import DepthwiseSeparableConv


class ZeroDCEPP(nn.Module):
    def __init__(self, scale_factor=4, num_features=32):
        super(ZeroDCEPP, self).__init__()

        self.relu = nn.ReLU(inplace=True)
        self.scale_factor = scale_factor
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=self.scale_factor)

        self.e_conv1 = DepthwiseSeparableConv(3, num_features) 
        self.e_conv2 = DepthwiseSeparableConv(num_features, num_features) 
        self.e_conv3 = DepthwiseSeparableConv(num_features, num_features) 
        self.e_conv4 = DepthwiseSeparableConv(num_features, num_features) 
        self.e_conv5 = DepthwiseSeparableConv(num_features * 2, num_features) # x3 + x4
        self.e_conv6 = DepthwiseSeparableConv(num_features * 2, num_features) # x2 + x5
        self.e_conv7 = DepthwiseSeparableConv(num_features * 2, 3)            # x1 + x6

    def apply_curve(self, x, x_r):
        for _ in range(8):
            x = x + x_r * (torch.pow(x, 2) - x)
        return x

    def forward(self, x):
        # 다운 스케일
        if self.scale_factor == 1:
            x_down = x
        else:
            x_down = F.interpolate(x, scale_factor=1/self.scale_factor, mode='bilinear', align_corners=False)


        # 특징 추출 레이어 시작 
        layer1 = self.relu(self.e_conv1(x_down))
        layer2 = self.relu(self.e_conv2(layer1))
        layer3 = self.relu(self.e_conv3(layer2))
        layer4 = self.relu(self.e_conv4(layer3))
        
        # Feature Concat
        layer5 = self.relu(self.e_conv5(torch.cat([layer3, layer4], 1)))
        layer6 = self.relu(self.e_conv6(torch.cat([layer2, layer5], 1)))
        layer7 = self.e_conv7(torch.cat([layer1, layer6], 1))
        

        x_r = torch.tanh(layer7)

        # 원본 사이즈로 다시 upsampling
        if self.scale_factor != 1:
            x_r = self.upsample(x_r)
        
        # 원본 이미지에 라이트맵을 곱함.
        enhance_image = self.apply_curve(x, x_r)
        
        return enhance_image, x_r

ONNX Model 배포를 위한 모델 코드

import torch


class OnnxModel(torch.nn.Module):
    def __init__(self, backbone:torch.nn.Module):
        super(OnnxModel, self).__init__()

        self.backbone = backbone


    def forward(self, x):
    
        ## 학습한 모델을 내보내기 할때 curve parameter는 inference시 필요가 없기 때문에. 
        ## OnnxModel용 클래스를 만들어서 curve를 제외한 결과만 리턴하도록 수정 
        
        x, curve = self.backbone(x)

        return x

Training Code

import torch
import torch.onnx
import numpy as np
import cv2
import os

from model.zerodcepp import ZeroDCEPP
from loss.zerodce_loss import ZeroDCETotalLoss
from dataset.lle_dataset import get_lle_loader
from model.onnx_model import OnnxModel

# GPU 사용 가능 여부 확인
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"현재 사용 중인 디바이스: {device}")

# Hyper Parameter
epochs = 50
learning_rate = 0.0001
weight_decay = 0.0001
batch_size = 8
scale_factor = 4
num_features = 32
image_width = 1024
image_height = 1024
image_channel = 3

dataset_path = "C://github//dataset//lol_dataset//our485//all"
weight_path = "C://github//Dot4Seminar//working//python//results//weights.pth"
onnx_model_path = "C://github//Dot4Seminar//working//python//results//model.onnx"


dummy_input = torch.randn(size=(1, image_channel, image_height, image_width)).to(device)

model = ZeroDCEPP(scale_factor=scale_factor, num_features=num_features)
model = model.to(device)
if os.path.exists(weight_path):
    state_dict = torch.load(weight_path, map_location=device)
    model.load_state_dict(state_dict)


dataloader = get_lle_loader(dataset_path, batch_size, resize_shape=(image_height, image_width))
total_batches = len(dataloader)


loss = ZeroDCETotalLoss().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

temp_loss = 1000000
for epoch in range(epochs):
    avg_loss = 0

    model.train()
    for i, x_image in enumerate(dataloader):

        gpu_x_image = x_image.to(device)
        enhanced_img, curve_params = model(gpu_x_image)
        current_loss = loss(gpu_x_image, enhanced_img, curve_params)

        optimizer.zero_grad()
        current_loss.backward()
        optimizer.step()

        avg_loss += current_loss.item() / total_batches

    tensor_input_check = gpu_x_image[0]
    tensor_output_check = enhanced_img[0]
    # 1. CPU 이동 및 Numpy 변환
    input = tensor_input_check.detach().cpu().numpy()
    input = np.transpose(input, (1, 2, 0))
    input = (input * 255).astype(np.uint8)
    input = cv2.cvtColor(input, cv2.COLOR_RGB2BGR)

    # 1. CPU 이동 및 Numpy 변환
    output = tensor_output_check.detach().cpu().numpy()
    output = np.transpose(output, (1, 2, 0))
    output = (output * 255).astype(np.uint8)
    output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)


    # 5. 시각화
    cv2.imshow("original", input)
    cv2.imshow("output", output)
    cv2.waitKey(1) # 1ms 대기 (학습 루프 멈춤 방지)

    if temp_loss > avg_loss:
        temp_loss = avg_loss
        
        model.eval()
    
        onnx_model = OnnxModel(backbone=model)

        torch.save(model.state_dict(), weight_path)
        torch.onnx.export(
            onnx_model,                      # 실행할 모델
            dummy_input,                # 모델 입력 예시
            onnx_model_path,            # 저장 파일명
            export_params=True,         # 모델 파일 안에 학습된 파라미터 저장
            opset_version=11,           # Bilinear 연산을 안정적으로 지원하는 버전
            do_constant_folding=True,   # 상수 폴딩 최적화 (속도 향상)
            input_names=['input'],      # 입력 노드 이름 (C++에서 호출 시 사용)
            output_names=['output'],    # 출력 노드 이름
        )
        
    print('current avg loss = ', avg_loss)

[파란매]

솔직히 알아듣지 못하지만 대단하시네요 제가 목표가 AI 는 모르지만 @gellston 님같이
일차로 학습을 시키는 방법을 배우는것데 대단하세요

[gellston]

음… 딥러닝 너무 어려운것 같습니다.

[gellston]

학습한 모델을 추론

지난 포스트에서 학습한 모델을 아래의 코드를 이용해서 추론해보겠습니다.
추론시에는 ONNX라이브러리를 사용하지 않고 pytorch 모델의 저장된 weights를 로드하여 추론 하였습니다.
코드 동작은 저조도 이미지가 저장되어 있는 이미지를 리스트로 불러와서 순회하며 추론을 진행하며 opencv를 이용해서 시각화 및 정상 동작하는지 파일로 저장하여 확인했습니다.

모델 추론 코드

import torch
import numpy as np
import cv2
import os


from model.zerodcepp import ZeroDCEPP


# GPU 사용 가능 여부 확인
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"현재 사용 중인 디바이스: {device}")



# Hyper Parameter
scale_factor = 4
num_features = 32
image_width = 1024
image_height = 1024
image_channel = 3

dataset_path = "C://github//dataset//lol_dataset//our485//low"
weight_path = "C://github//Dot4Seminar//working//python//results//weights.pth"
output_dir = "C://github//Dot4Seminar//working//python//results//output_images"


dummy_input = torch.randn(size=(1, image_channel, image_height, image_width)).to(device)

model = ZeroDCEPP(scale_factor=scale_factor, num_features=num_features)
model = model.to(device)
if os.path.exists(weight_path):
    state_dict = torch.load(weight_path, map_location=device)
    model.load_state_dict(state_dict)

model.eval()

os.makedirs(output_dir, exist_ok=True)

file_list = os.listdir(dataset_path)
image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')


for file_name in file_list:
    if not file_name.lower().endswith(image_extensions):
        continue

    full_path = os.path.join(dataset_path, file_name)
    input = cv2.imread(full_path)
    
    if input is None:
        continue

    h, w, _ = input.shape

    input_img = cv2.cvtColor(input, cv2.COLOR_BGR2RGB)
    input_img = cv2.resize(input_img, (image_width, image_height))
    input_img = input_img.astype(np.float32) / 255.0
    input_img = np.transpose(input_img, (2, 0, 1))
    
    input_tensor = torch.from_numpy(input_img).unsqueeze(0).to(device)

    with torch.no_grad():
        enhanced_tensor, _ = model(input_tensor)


    output = enhanced_tensor.squeeze(0).cpu().numpy()
    output = np.transpose(output, (1, 2, 0))
    output = (output * 255.0).clip(0, 255).astype(np.uint8)
    
    output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
    output = cv2.resize(output, (w, h))


    save_path = os.path.join(output_dir, file_name)
    cv2.imwrite(save_path, output)


    cv2.imshow('input', input)
    cv2.imshow('output', output)
    

    cv2.waitKey(1)

시연 동영상

click image to open video

결과 이미지

#1 원본

#1 결과

#2 원본

#2 결과

#3 원본

#3 결과

일부 저조도가 아닌 밝은 이미지에서는 왜 오동작을 할가?

왜냐면 ZeroDCE의 학습 로스는 어두운 이미지를 밝게 만들도록 weights 들을 조정하지만 (curve parameters) 반대로 과하게 밝은 픽셀 부분을 조정하여 낮추는 설계는 들어가 있지 않아서 그런것 같습니다.
만약에 이 모델을 사용해야 한다면 평균 밝기가 어느 값 이상일때에는 사용하지 않도록 (굳이 저조도 상태가 아니기떄문에?) 예외처리를 한다거나 아니면 학습중에 설정하는 하이퍼파라미터를 조정하여 타겟 밝기를 너무 밝지 않도록 조정하는 방법이 있을 수 있겠네요 (해봤는데 잘안됩니다..)

[windowmini]

우오….. CCTV 스크린샷 이 어두울때 사용 하면… 대박 나는거 아님까 ?

음 대기업만 쓰려나

능력자 횽아네… .. 논문을 개발로…

[tkm]

생각보다 코드가 간단해서 onnx 변환만 되면 쉽게 사용할 수 있겠네요

[BigSquare]

@gellston

반대로 과하게 밝은 픽셀 부분을 조정하여 낮추는 설계는 들어가 있지 않아서 그런것 같습니다.

평균 밝기가 일정 이상일 때,

이미지 반전 → 로직 적용 → 결과물 반전

으로 한 번 해보는 것도 좋으리라 생각됩니다.

[gellston]

요즘 CCTV에 딥러닝 쓴다고 들었는데 더 좋은게 들어가 있지 않을가요?ㅠ 잘모르겠군요 거기 업계는 ㅠㅠ

[gellston]

네 맞습니다 논문 저자분이 간결하게 설계를 하셨더라구요
오래된 논문이지만 저 저도 논문에서 항상 reference로 나오는 논문같더군요.
onnx모델은 누겟으로 감싸서 만들때 같이 readme를 작성해서 같이 공유드리겠습니다 (공부용)

[gellston]

아 좋은 아이디어 같습니다.
C++ 래핑을 할건데 그때 worst케이스 이미지들을 모아서 되는지 실험해봐야겠군요.

대충 이런 느낌이군요?
다만 딥러닝이 저 반전된 이미지를 입력으로 받을때 어떻게 해석할지 (오동작) 해봐야 될 것 같네요
아마 잘 되리라 보는데… 만약에 반전이 안된다면 학습할 때 augmentation으로 픽셀 값 반전을 넣고 학습한 모델을 가지고 아이디어를 적용해보겠습니다.

의견 주셔서 감사합니다!

[gellston]

아까 의견 주신 거 댓글로 답변 드리다가 그러니까 저 그림을 그리다가?
문득 드는 생각이 이미 로스 함수에서는 인접 픽셀의 밝기차가(국소적인 밝기차) 과하지 않도록 제약을 주고있는데요.
거기에 더해서 평균 밝기가 내가 설정한 밝기로 따라가도록 설계가 되어있는데.
실험에서 입력으로 넣은 정상적인 이미지는 사람의 눈으로 보았을때 전체적인 글로벌한 밝기가
올라가 있는 상태인데 모델의 피쳐를 뽑는 컨볼루션 레이어들의 커널사이즈는 3x3으로 작기 때문에 전체 씬 해석 능력이 떨어지는 것도 어느부분 작용했을 거라고 생각이 드는군요 ㅠ
모델이 너무 좁게 보는게 아닌가 생각이 음..
그래서 말씀하신 반전 실험에 더해서 피쳐레이어에 커널 사이즈를 늘려서 테스트해보는 것도 해보도록 하겠습니다.

감사합니다!