import os
import numpy as np
os.environ["MKL_NUM_THREADS"] = "1"  # noqa F402
os.environ["NUMEXPR_NUM_THREADS"] = "1"  # noqa F402
os.environ["OMP_NUM_THREADS"] = "1"  # noqa F402

from .mono_dataset import MonoDataset


class WAYMO(MonoDataset):
    """Superclass for different types of Waymo dataset loader
    """
    def __init__(self, *args, **kwargs):
        super(WAYMO, self).__init__(*args, **kwargs)
        self.K = np.array([[1.07, 0, 0.5, 0],
                        [0, 1.61, 0.5, 0],
                        [0, 0, 1, 0],
                        [0, 0, 0, 1]], dtype=np.float32)
        

    def get_color(self, line):
        color = self.loader(os.path.join(self.data_path, line))

        return color
