# Path: third_party/data/base_dataset.py
# class BaseDataset(data.Dataset, ABC):
#     """This class is an abstract base class (ABC) for datasets.
# 
#     To create a subclass, you need to implement the following four functions:
#     -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
#     -- <__len__>:                       return the size of dataset.
#     -- <__getitem__>:                   get a data point.
#     -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
#     """
# 
#     def __init__(self, opt):
#         """Initialize the class; save the options in the class
# 
#         Parameters:
#             opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
#         """
#         self.opt = opt
#         self.root = opt.dataroot
# 
#     @staticmethod
#     def modify_commandline_options(parser, is_train):
#         """Add new dataset-specific options, and rewrite default values for existing options.
# 
#         Parameters:
#             parser          -- original option parser
#             is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
# 
#         Returns:
#             the modified parser.
#         """
#         return parser
# 
#     @abstractmethod
#     def __len__(self):
#         """Return the total number of images in the dataset."""
#         return 0
# 
#     @abstractmethod
#     def __getitem__(self, index):
#         """Return a data point and its metadata information.
# 
#         Parameters:
#             index - - a random integer for data indexing
# 
#         Returns:
#             a dictionary of data with their names. It ususally contains the data itself and its metadata information.
#         """
#         pass
# 
# Path: third_party/data/image_folder.py
# def make_dataset(dir, max_dataset_size=float("inf")):
#     images = []
#     assert os.path.isdir(dir), '%s is not a valid directory' % dir
# 
#     for root, _, fnames in sorted(os.walk(dir)):
#         for fname in fnames:
#             if is_image_file(fname):
#                 path = os.path.join(root, fname)
#                 images.append(path)
#     images = sorted(images)
#     return images[:min(max_dataset_size, len(images))]


import cv2
import torchvision.transforms as transforms
import torch.nn.functional as F
import os
import torch
import numpy as np
import json
from third_party.data.base_dataset import BaseDataset
from third_party.data.image_folder import make_dataset
from PIL import Image

    It assumes that the directory specified by 'dataroot' contains metadata.json, and the directories iuv, rgb_256, and rgb_512.
    The 'iuv' directory should contain directories named 01, 02, etc. for each layer, each containing per-frame UV images.
    """
    @staticmethod
    def modify_commandline_options(parser, is_train):
        parser.add_argument('--height', type=int, default=256, help='image height')
        parser.add_argument('--width', type=int, default=448, help='image width')
        parser.add_argument('--trimap_width', type=int, default=20, help='trimap gray area width')
        parser.add_argument('--use_mask_images', action='store_true', default=False, help='use custom masks')
        parser.add_argument('--use_homographies', action='store_true', default=False, help='handle camera motion')
        return parser

    def __init__(self, opt):
        """Initialize this dataset class.

        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseDataset.__init__(self, opt)
        rgbdir = os.path.join(opt.dataroot, 'rgb_256')
        if opt.do_upsampling:
            rgbdir = os.path.join(opt.dataroot, 'rgb_512')
        uvdir = os.path.join(opt.dataroot, 'iuv')
        self.image_paths = sorted(make_dataset(rgbdir, opt.max_dataset_size))
        n_images = len(self.image_paths)
        layers = sorted(os.listdir(uvdir))
        layers = [l for l in layers if l.isdigit()]
        self.iuv_paths = []
        for l in layers:

# Path: third_party/data/base_dataset.py
# class BaseDataset(data.Dataset, ABC):
#     """This class is an abstract base class (ABC) for datasets.
# 
#     To create a subclass, you need to implement the following four functions:
#     -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
#     -- <__len__>:                       return the size of dataset.
#     -- <__getitem__>:                   get a data point.
#     -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
#     """
# 
#     def __init__(self, opt):
#         """Initialize the class; save the options in the class
# 
#         Parameters:
#             opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
#         """
#         self.opt = opt
#         self.root = opt.dataroot
# 
#     @staticmethod
#     def modify_commandline_options(parser, is_train):
#         """Add new dataset-specific options, and rewrite default values for existing options.
# 
#         Parameters:
#             parser          -- original option parser
#             is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
# 
#         Returns:
#             the modified parser.
#         """
#         return parser
# 
#     @abstractmethod
#     def __len__(self):
#         """Return the total number of images in the dataset."""
#         return 0
# 
#     @abstractmethod
#     def __getitem__(self, index):
#         """Return a data point and its metadata information.
# 
#         Parameters:
#             index - - a random integer for data indexing
# 
#         Returns:
#             a dictionary of data with their names. It ususally contains the data itself and its metadata information.
#         """
#         pass
# 
# Path: third_party/data/image_folder.py
# def make_dataset(dir, max_dataset_size=float("inf")):
#     images = []
#     assert os.path.isdir(dir), '%s is not a valid directory' % dir
# 
#     for root, _, fnames in sorted(os.walk(dir)):
#         for fname in fnames:
#             if is_image_file(fname):
#                 path = os.path.join(root, fname)
#                 images.append(path)
#     images = sorted(images)
#     return images[:min(max_dataset_size, len(images))]


# Path: data/layered_video_dataset.py
import cv2
import torchvision.transforms as transforms
import torch.nn.functional as F
import os
import torch
import numpy as np
import json
from third_party.data.base_dataset import BaseDataset
from third_party.data.image_folder import make_dataset
from PIL import Image


    It assumes that the directory specified by 'dataroot' contains metadata.json, and the directories iuv, rgb_256, and rgb_512.
    The 'iuv' directory should contain directories named 01, 02, etc. for each layer, each containing per-frame UV images.
    """
    @staticmethod
    def modify_commandline_options(parser, is_train):
        parser.add_argument('--height', type=int, default=256, help='image height')
        parser.add_argument('--width', type=int, default=448, help='image width')
        parser.add_argument('--trimap_width', type=int, default=20, help='trimap gray area width')
        parser.add_argument('--use_mask_images', action='store_true', default=False, help='use custom masks')
        parser.add_argument('--use_homographies', action='store_true', default=False, help='handle camera motion')
        return parser

    def __init__(self, opt):
        """Initialize this dataset class.

        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseDataset.__init__(self, opt)
        rgbdir = os.path.join(opt.dataroot, 'rgb_256')
        if opt.do_upsampling:
            rgbdir = os.path.join(opt.dataroot, 'rgb_512')
        uvdir = os.path.join(opt.dataroot, 'iuv')
        self.image_paths = sorted(make_dataset(rgbdir, opt.max_dataset_size))
        n_images = len(self.image_paths)
        layers = sorted(os.listdir(uvdir))
        layers = [l for l in layers if l.isdigit()]
        self.iuv_paths = []
        for l in layers:

            layer_iuv_paths = sorted(make_dataset(os.path.join(uvdir, l), n_images))