EMINIME/URetinex-Net

1import torch2import torch.nn as nn3from torchvision.transforms import Grayscale4 5 6class P(nn.Module):7    """8        to solve min(P) = ||I-PQ||^2 + γ||P-R||^29        this is a least square problem10        how to solve?11        P* = (gamma*R + I*Q) / (Q*Q + gamma)12    """13    def __init__(self):14        super().__init__()15 16    def forward(self, I, Q, R, gamma):17        return ((I * Q + gamma * R) / (gamma + Q * Q))18        19class Q(nn.Module):20    """21        to solve min(Q) = ||I-PQ||^2 + λ||Q-L||^222        Q* = (lamda*L + I*P) / (P*P + lamda)23    """24    def __init__(self):25        super().__init__()26    27    def forward(self, I, P, L, lamda):28        29        IR = I[:, 0:1, :, :]30        IG = I[:, 1:2, :, :]31        IB = I[:, 2:3, :, :]32 33        PR = P[:, 0:1, :, :]34        PG = P[:, 1:2, :, :]35        PB = P[:, 2:3, :, :]36 37        return (IR*PR + IG*PG + IB*PB + lamda*L) / ((PR*PR + PG*PG + PB*PB) + lamda)38