From 6169d3daac1d7ab4c5ba926495c19cff23c12d74 Mon Sep 17 00:00:00 2001 From: vlpr Date: Wed, 15 May 2024 10:59:56 +0000 Subject: [PATCH] minor name change --- src/models/srlut.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/models/srlut.py b/src/models/srlut.py index ed80dda..11bee80 100644 --- a/src/models/srlut.py +++ b/src/models/srlut.py @@ -4,7 +4,7 @@ import torch.nn.functional as F import numpy as np from pathlib import Path from common.lut import forward_2x2_input_SxS_output -from common.layers import RgbToYcbcr, YcbcrToRgb +from common import layers class SRLut(nn.Module): def __init__( @@ -39,13 +39,13 @@ class SRLut(nn.Module): -class SRLutRot90(nn.Module): +class SRLutR90(nn.Module): def __init__( self, quantization_interval, scale ): - super(SRLutRot90, self).__init__() + super(SRLutR90, self).__init__() self.scale = scale self.quantization_interval = quantization_interval self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) @@ -77,18 +77,18 @@ class SRLutRot90(nn.Module): return f"{self.__class__.__name__}\n lut size: {self.stage_lut.shape}" -class SRLutRot90Y(nn.Module): +class SRLutR90Y(nn.Module): def __init__( self, quantization_interval, scale ): - super(SRLutRot90Y, self).__init__() + super(SRLutR90Y, self).__init__() self.scale = scale self.quantization_interval = quantization_interval self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32)) - self.rgb_to_ycbcr = RgbToYcbcr() - self.ycbcr_to_rgb = YcbcrToRgb() + self.rgb_to_ycbcr = layers.RgbToYcbcr() + self.ycbcr_to_rgb = layers.YcbcrToRgb() @staticmethod def init_from_lut(