MLBench/ReaLens

1from .pix2pix_model import Pix2PixModel2import torch3from skimage import color  # used for lab2rgb4import numpy as np5 6 7class ColorizationModel(Pix2PixModel):8    """This is a subclass of Pix2PixModel for image colorization (black & white image -> colorful images).9 10    The model training requires '-dataset_model colorization' dataset.11    It trains a pix2pix model, mapping from L channel to ab channels in Lab color space.12    By default, the colorization dataset will automatically set '--input_nc 1' and '--output_nc 2'.13    """14 15    @staticmethod16    def modify_commandline_options(parser, is_train=True):17        """Add new dataset-specific options, and rewrite default values for existing options.18 19        Parameters:20            parser          -- original option parser21            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.22 23        Returns:24            the modified parser.25 26        By default, we use 'colorization' dataset for this model.27        See the original pix2pix paper (https://arxiv.org/pdf/1611.07004.pdf) and colorization results (Figure 9 in the paper)28        """29        Pix2PixModel.modify_commandline_options(parser, is_train)30        parser.set_defaults(dataset_mode="colorization")31        return parser32 33    def __init__(self, opt):34        """Initialize the class.35 36        Parameters:37            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions38 39        For visualization, we set 'visual_names' as 'real_A' (input real image),40        'real_B_rgb' (ground truth RGB image), and 'fake_B_rgb' (predicted RGB image)41        We convert the Lab image 'real_B' (inherited from Pix2pixModel) to a RGB image 'real_B_rgb'.42        we convert the Lab image 'fake_B' (inherited from Pix2pixModel) to a RGB image 'fake_B_rgb'.43        """44        # reuse the pix2pix model45        Pix2PixModel.__init__(self, opt)46        # specify the images to be visualized.47        self.visual_names = ["real_A", "real_B_rgb", "fake_B_rgb"]48 49    def lab2rgb(self, L, AB):50        """Convert an Lab tensor image to a RGB numpy output51        Parameters:52            L  (1-channel tensor array): L channel images (range: [-1, 1], torch tensor array)53            AB (2-channel tensor array):  ab channel images (range: [-1, 1], torch tensor array)54 55        Returns:56            rgb (RGB numpy image): rgb output images  (range: [0, 255], numpy array)57        """58        AB2 = AB * 110.059        L2 = (L + 1.0) * 50.060        Lab = torch.cat([L2, AB2], dim=1)61        Lab = Lab[0].data.cpu().float().numpy()62        Lab = np.transpose(Lab.astype(np.float64), (1, 2, 0))63        rgb = color.lab2rgb(Lab) * 25564        return rgb65 66    def compute_visuals(self):67        """Calculate additional output images for visdom and HTML visualization"""68        self.real_B_rgb = self.lab2rgb(self.real_A, self.real_B)69        self.fake_B_rgb = self.lab2rgb(self.real_A, self.fake_B)70