45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
import torch
|
|
from torch import nn
|
|
import torch.nn.functional as F
|
|
import torchvision.transforms.functional as tf
|
|
|
|
from .filter import Filter
|
|
from .backbone import EfficientBackbone
|
|
from .module import CascadeArgumentRegressor, FilterPerformer
|
|
|
|
|
|
class Harmonizer(nn.Module):
|
|
def __init__(self):
|
|
super(Harmonizer, self).__init__()
|
|
|
|
self.input_size = (256, 256)
|
|
self.filter_types = [
|
|
Filter.TEMPERATURE,
|
|
Filter.BRIGHTNESS,
|
|
Filter.CONTRAST,
|
|
Filter.SATURATION,
|
|
Filter.HIGHLIGHT,
|
|
Filter.SHADOW,
|
|
]
|
|
|
|
self.backbone = EfficientBackbone.from_name('efficientnet-b0')
|
|
self.regressor = CascadeArgumentRegressor(1280, 160, 1, len(self.filter_types))
|
|
self.performer = FilterPerformer(self.filter_types)
|
|
|
|
def predict_arguments(self, comp, mask):
|
|
comp = F.interpolate(comp, self.input_size, mode='bilinear', align_corners=False)
|
|
mask = F.interpolate(mask, self.input_size, mode='bilinear', align_corners=False)
|
|
|
|
fg = torch.cat((comp, mask), dim=1)
|
|
bg = torch.cat((comp, (1 - mask)), dim=1)
|
|
|
|
enc2x, enc4x, enc8x, enc16x, enc32x = self.backbone(fg, bg)
|
|
arguments = self.regressor(enc32x)
|
|
return arguments
|
|
|
|
def restore_image(self, comp, mask, arguments):
|
|
assert len(arguments) == len(self.filter_types)
|
|
|
|
arguments = [torch.clamp(arg, -1, 1).view(-1, 1, 1, 1) for arg in arguments]
|
|
return self.performer.restore(comp, mask, arguments)
|