from argparse import ArgumentParser, Namespace
from ctypes import Union

from datasets.compare_base import CompareBase
from ml_lib.utils.tools import add_argparse_args


class PrimatesLibrosaDatamodule(CompareBase):

    class_names = ['background', 'chimpanze', 'geunon', 'mandrille', 'redcap']
    sub_dataset_name = 'primates'

    def __init__(self, *args, **kwargs):
        super(PrimatesLibrosaDatamodule, self).__init__(*args, **kwargs)

    @classmethod
    def add_argparse_args(cls, parent_parser):
        return add_argparse_args(CompareBase, parent_parser)


    @classmethod
    def from_argparse_args(cls, args, **kwargs):
        return CompareBase.from_argparse_args(args, class_names=cls.class_names, sub_dataset_name=cls.sub_dataset_name)