Frank Schwenk
a60a18a253
Introduce ruff lint/format config, expand .gitignore, and reformat Python sources. Co-authored-by: Cursor <cursoragent@cursor.com>
251 lines
7.9 KiB
Python
251 lines
7.9 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from imagepipeline.ai.hdrnet.slice import batch_bilateral_slice
|
|
|
|
|
|
class ConvBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
inc,
|
|
outc,
|
|
kernel_size=3,
|
|
padding=1,
|
|
stride=1,
|
|
use_bias=True,
|
|
activation=nn.ReLU,
|
|
batch_norm=False,
|
|
) -> None:
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(
|
|
int(inc), int(outc), kernel_size, padding=padding, stride=stride, bias=use_bias
|
|
)
|
|
self.activation = activation() if activation else None
|
|
self.bn = nn.BatchNorm2d(outc) if batch_norm else None
|
|
if use_bias and not batch_norm:
|
|
self.conv.bias.data.fill_(0.0)
|
|
torch.nn.init.kaiming_uniform_(self.conv.weight)
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
if self.bn is not None:
|
|
x = self.bn(x)
|
|
if self.activation is not None:
|
|
x = self.activation(x)
|
|
return x
|
|
|
|
|
|
class FC(nn.Module):
|
|
def __init__(self, inc, outc, activation=nn.ReLU, batch_norm=False) -> None:
|
|
super().__init__()
|
|
self.fc = nn.Linear(int(inc), int(outc), bias=(not batch_norm))
|
|
self.activation = activation() if activation else None
|
|
self.bn = nn.BatchNorm1d(outc) if batch_norm else None
|
|
if not batch_norm:
|
|
self.fc.bias.data.fill_(0.0)
|
|
torch.nn.init.kaiming_uniform_(self.fc.weight)
|
|
|
|
def forward(self, x):
|
|
x = self.fc(x)
|
|
if self.bn is not None:
|
|
x = self.bn(x)
|
|
if self.activation is not None:
|
|
x = self.activation(x)
|
|
return x
|
|
|
|
|
|
class Slice(nn.Module):
|
|
def forward(self, bilateral_grid, guidemap):
|
|
bilateral_grid = bilateral_grid.permute(0, 3, 4, 2, 1)
|
|
guidemap = guidemap.squeeze(1)
|
|
coeffs = batch_bilateral_slice(bilateral_grid, guidemap).permute(0, 3, 1, 2)
|
|
return coeffs
|
|
|
|
|
|
class ApplyCoeffs(nn.Module):
|
|
def forward(self, coeff, full_res_input):
|
|
r = (
|
|
torch.sum(full_res_input * coeff[:, 0:3, :, :], dim=1, keepdim=True)
|
|
+ coeff[:, 9:10, :, :]
|
|
)
|
|
g = (
|
|
torch.sum(full_res_input * coeff[:, 3:6, :, :], dim=1, keepdim=True)
|
|
+ coeff[:, 10:11, :, :]
|
|
)
|
|
b = (
|
|
torch.sum(full_res_input * coeff[:, 6:9, :, :], dim=1, keepdim=True)
|
|
+ coeff[:, 11:12, :, :]
|
|
)
|
|
return torch.cat([r, g, b], dim=1)
|
|
|
|
|
|
class GuideNN(nn.Module):
|
|
def __init__(self, params) -> None:
|
|
super().__init__()
|
|
self.conv1 = ConvBlock(
|
|
3, params["guide_complexity"], kernel_size=1, padding=0, batch_norm=True
|
|
)
|
|
self.conv2 = ConvBlock(
|
|
params["guide_complexity"], 1, kernel_size=1, padding=0, activation=nn.Sigmoid
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.conv2(self.conv1(x))
|
|
|
|
|
|
class Coeffs(nn.Module):
|
|
def __init__(self, nin=4, nout=3, params=None) -> None:
|
|
super().__init__()
|
|
self.params = params
|
|
self.nin = nin
|
|
self.nout = nout
|
|
lb = params["luma_bins"]
|
|
cm = params["channel_multiplier"]
|
|
sb = params["spatial_bin"]
|
|
bn = params["batch_norm"]
|
|
nsize = params["net_input_size"]
|
|
|
|
n_layers_splat = int(np.log2(nsize / sb))
|
|
self.splat_features = nn.ModuleList()
|
|
prev_ch = 3
|
|
for index in range(n_layers_splat):
|
|
use_bn = bn if index > 0 else False
|
|
out_ch = cm * (2**index) * lb
|
|
self.splat_features.append(ConvBlock(prev_ch, out_ch, 3, stride=2, batch_norm=use_bn))
|
|
prev_ch = out_ch
|
|
splat_ch = prev_ch
|
|
|
|
n_layers_global = int(np.log2(sb / 4))
|
|
self.global_features_conv = nn.ModuleList()
|
|
self.global_features_fc = nn.ModuleList()
|
|
for _ in range(n_layers_global):
|
|
self.global_features_conv.append(
|
|
ConvBlock(prev_ch, cm * 8 * lb, 3, stride=2, batch_norm=bn)
|
|
)
|
|
prev_ch = cm * 8 * lb
|
|
|
|
n_total = n_layers_splat + n_layers_global
|
|
prev_ch = int(prev_ch * (nsize / 2**n_total) ** 2)
|
|
self.global_features_fc.append(FC(prev_ch, 32 * cm * lb, batch_norm=bn))
|
|
self.global_features_fc.append(FC(32 * cm * lb, 16 * cm * lb, batch_norm=bn))
|
|
self.global_features_fc.append(
|
|
FC(16 * cm * lb, 8 * cm * lb, activation=None, batch_norm=bn)
|
|
)
|
|
|
|
self.local_features = nn.ModuleList(
|
|
[
|
|
ConvBlock(splat_ch, 8 * cm * lb, 3, batch_norm=bn),
|
|
ConvBlock(8 * cm * lb, 8 * cm * lb, 3, activation=None, use_bias=False),
|
|
]
|
|
)
|
|
self.conv_out = ConvBlock(8 * cm * lb, lb * nout * nin, 1, padding=0, activation=None)
|
|
self.relu = nn.ReLU()
|
|
|
|
def forward(self, lowres_input):
|
|
params = self.params
|
|
bs = lowres_input.shape[0]
|
|
lb = params["luma_bins"]
|
|
cm = params["channel_multiplier"]
|
|
|
|
x = lowres_input
|
|
for layer in self.splat_features:
|
|
x = layer(x)
|
|
splat_features = x
|
|
|
|
for layer in self.global_features_conv:
|
|
x = layer(x)
|
|
x = x.view(bs, -1)
|
|
for layer in self.global_features_fc:
|
|
x = layer(x)
|
|
global_features = x
|
|
|
|
x = splat_features
|
|
for layer in self.local_features:
|
|
x = layer(x)
|
|
fusion = self.relu(x + global_features.view(bs, 8 * cm * lb, 1, 1))
|
|
x = self.conv_out(fusion)
|
|
return torch.stack(torch.split(x, self.nin * self.nout, 1), 2)
|
|
|
|
|
|
class HDRPointwiseNN(nn.Module):
|
|
def __init__(self, params) -> None:
|
|
super().__init__()
|
|
self.coeffs = Coeffs(params=params)
|
|
self.guide = GuideNN(params=params)
|
|
self.slice = Slice()
|
|
self.apply_coeffs = ApplyCoeffs()
|
|
|
|
def forward(self, lowres, fullres):
|
|
coeffs = self.coeffs(lowres)
|
|
guide = self.guide(fullres)
|
|
slice_coeffs = self.slice(coeffs, guide)
|
|
return self.apply_coeffs(slice_coeffs, fullres)
|
|
|
|
|
|
def default_hdrnet_params(net_input_size: int = 256) -> dict:
|
|
return {
|
|
"luma_bins": 8,
|
|
"channel_multiplier": 1,
|
|
"spatial_bin": 16,
|
|
"batch_norm": True,
|
|
"net_input_size": net_input_size,
|
|
"guide_complexity": 16,
|
|
}
|
|
|
|
|
|
def load_hdrnet_checkpoint(checkpoint_path, device: torch.device):
|
|
state = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
|
if "model_params" in state:
|
|
params = state["model_params"]
|
|
del state["model_params"]
|
|
else:
|
|
params = default_hdrnet_params()
|
|
model = HDRPointwiseNN(params=params)
|
|
model.load_state_dict(state)
|
|
model.to(device)
|
|
model.eval()
|
|
return model, params
|
|
|
|
|
|
def resize_rgb_array(arr: np.ndarray, size: int) -> np.ndarray:
|
|
from PIL import Image
|
|
|
|
image = Image.fromarray(arr.astype(np.uint8))
|
|
short = min(image.size)
|
|
scale = size / short
|
|
new_size = (max(1, round(image.size[0] * scale)), max(1, round(image.size[1] * scale)))
|
|
return np.asarray(image.resize(new_size, Image.Resampling.NEAREST))
|
|
|
|
|
|
def enhance_image_hdrnet(
|
|
model: HDRPointwiseNN,
|
|
image_rgb,
|
|
*,
|
|
device: torch.device,
|
|
net_input_size: int,
|
|
strength: float = 1.0,
|
|
):
|
|
from PIL import Image
|
|
|
|
if not isinstance(image_rgb, Image.Image):
|
|
image_rgb = Image.fromarray(image_rgb)
|
|
|
|
full_arr = np.asarray(image_rgb, dtype=np.float32)
|
|
low_arr = resize_rgb_array(full_arr, net_input_size)
|
|
low = torch.from_numpy(low_arr).permute(2, 0, 1).unsqueeze(0).float() / 255.0
|
|
full = torch.from_numpy(full_arr).permute(2, 0, 1).unsqueeze(0).float() / 255.0
|
|
low = low.to(device)
|
|
full = full.to(device)
|
|
|
|
with torch.no_grad():
|
|
out = model(low, full)
|
|
if strength < 1.0:
|
|
out = full * (1.0 - strength) + out * strength
|
|
out = torch.clamp(out, 0.0, 1.0)
|
|
|
|
result = (out.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0).astype(np.uint8)
|
|
return Image.fromarray(result)
|