"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import time
from PIL import Image
from lavis.common.utils import get_abs_path, get_cache_path
from multiprocessing import Pool
from omegaconf import OmegaConf
from pathlib import Path
from torchvision.transforms import functional as TF
from tqdm import tqdm
import glob
import io
import json
import magic  # pip install python-magic
import numpy as np
import os
import pandas as pd
import requests
import shelve
import zlib

headers = {
    #'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/71.0.3578.98 Safari/537.36',
    "User-Agent": "Googlebot-Image/1.0",  # Pretend to be googlebot
    "X-Forwarded-For": "64.18.15.200",
}

def _df_split_apply(tup_arg):
    split_ind, subset, func = tup_arg
    r = subset.apply(func, axis=1)
    return (split_ind, r)

def df_multiprocess(df, processes, chunk_size, func, dataset_name):
    print("Generating parts...")
    with shelve.open(
        "%s_%s_%s_results.tmp" % (dataset_name, func.__name__, chunk_size)
    ) as results:

        pbar = tqdm(total=len(df), position=0)
        # Resume:
        finished_chunks = set([int(k) for k in results.keys()])
        pbar.desc = "Resuming"
        for k in results.keys():
            pbar.update(len(results[str(k)][1]))

        pool_data = (
            (index, df[i : i + chunk_size], func)
            for index, i in enumerate(range(0, len(df), chunk_size))
            if index not in finished_chunks
        )
        print(
            int(len(df) / chunk_size),
            "parts.",
            chunk_size,
            "per part.",
            "Using",
            processes,
            "processes",
        )

        pbar.desc = "Downloading"
        with Pool(processes) as pool:
            for i, result in enumerate(
                pool.imap_unordered(_df_split_apply, pool_data, 2)
            ):
                results[str(result[0])] = result
                pbar.update(len(result[1]))
        pbar.close()

    print("Finished Downloading.")
    return

# Unique name based on url
def _file_name(row):
    name = (
        "%s/%s_%s"
        % (
            # row["folder"],
            storage_dir,
            row.name,
            (zlib.crc32(row["url"].encode("utf-8")) & 0xFFFFFFFF),
        )
        + ".jpg"
    )
    return name

# For checking mimetypes separately without download
def check_mimetype(row):
    if os.path.isfile(str(row["file"])):
        row["mimetype"] = magic.from_file(row["file"], mime=True)
        row["size"] = os.stat(row["file"]).st_size
    return row

# Don't download image, just check with a HEAD request, can't resume.
# Can use this instead of download_image to get HTTP status codes.
def check_download(row):
    fname = _file_name(row)
    try:
        # not all sites will support HEAD
        response = requests.head(
            row["url"], stream=False, timeout=5, allow_redirects=True, headers=headers
        )
        row["status"] = response.status_code
        row["headers"] = dict(response.headers)
    except:
        # log errors later, set error as 408 timeout
        row["status"] = 408
        return row
    if response.ok:
        row["file"] = fname
    return row

def resize_img(req):
    image = Image.open(req).convert("RGB")
    image = TF.resize(
        # image, size=(resize_size, resize_size)
        image,
        size=resize_size,
    )  # , interpolation=Image.LANCZOS)
    return image

def download_image(row):
    fname = _file_name(row)
    # Skip Already downloaded, retry others later
    if os.path.isfile(fname):
        row["status"] = 200
        row["file"] = fname
        row["mimetype"] = magic.from_file(row["file"], mime=True)
        row["size"] = os.stat(row["file"]).st_size
        return row

    try:
        # use smaller timeout to skip errors, but can result in failed downloads
        response = requests.get(
            row["url"], stream=False, timeout=5, allow_redirects=True, headers=headers
        )
        row["status"] = response.status_code
        # row['headers'] = dict(response.headers)
    except Exception as e:
        # log errors later, set error as 408 timeout
        row["status"] = 408
        return row

    if response.ok:
        try:
            # some sites respond with gzip transport encoding
            response.raw.decode_content = True
            img = resize_img(io.BytesIO(response.content))
            img.save(fname)

            row["mimetype"] = magic.from_file(fname, mime=True)
            row["size"] = os.stat(fname).st_size

        except Exception as e:
            #     # This is if it times out during a download or decode
            row["status"] = 408

    row["file"] = fname
    return row

def open_tsv(fname, folder):
    print("Opening %s Data File..." % fname)
    df = pd.read_csv(
        fname, sep="\t", names=["url", "caption"]
    )  # , usecols=range(1, 2))
    df["folder"] = folder
    print("Processing", len(df), " Images:")
    return df

def df_from_shelve(chunk_size, func, dataset_name):
    print("Generating Dataframe from results...")
    with shelve.open(
        "%s_%s_%s_results.tmp" % (dataset_name, func.__name__, chunk_size)
    ) as results:
        keylist = sorted([int(k) for k in results.keys()])
        df = pd.concat([results[str(k)][1] for k in keylist], sort=True)
    return df

resize_size = 384

config_path = get_abs_path("configs/datasets/conceptual_caption/defaults_12m.yaml")

storage_dir = OmegaConf.load(
    config_path
).datasets.conceptual_caption_12m.build_info.images.storage
storage_dir = Path(get_cache_path(storage_dir))

os.makedirs(storage_dir, exist_ok=True)

# number of processes in the pool can be larger than cores
num_processes = 96
# num_processes = 1
# chunk_size is how many images per chunk per process - changing this resets progress when restarting.
images_per_part = 100

data_name = "cc12m"
# os.makedirs(data_name, exist_ok=True)

df = open_tsv("cc12m.tsv", data_name)
df_multiprocess(
    df=df,
    processes=num_processes,
    chunk_size=images_per_part,
    func=download_image,
    dataset_name=data_name,
)
df = df_from_shelve(
    chunk_size=images_per_part, func=download_image, dataset_name=data_name
)
df.to_csv(
    "downloaded_%s_report.tsv.gz" % data_name,
    compression="gzip",
    sep="\t",
    header=False,
    index=False,
)
print("Saved.")


from torchvision.transforms import functional as TF
from tqdm import tqdm
import glob
import io
import json
import magic  # pip install python-magic
import numpy as np
import os
import pandas as pd
import requests
import shelve
import zlib

headers = {
    #'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/71.0.3578.98 Safari/537.36',
    "User-Agent": "Googlebot-Image/1.0",  # Pretend to be googlebot
    "X-Forwarded-For": "64.18.15.200",
}

def _df_split_apply(tup_arg):
    split_ind, subset, func = tup_arg
    r = subset.apply(func, axis=1)
    return (split_ind, r)

def df_multiprocess(df, processes, chunk_size, func, dataset_name):
    print("Generating parts...")
    with shelve.open(
        "%s_%s_%s_results.tmp" % (dataset_name, func.__name__, chunk_size)
    ) as results:

        pbar = tqdm(total=len(df), position=0)
        # Resume:
        finished_chunks = set([int(k) for k in results.keys()])
        pbar.desc = "Resuming"
        for k in results.keys():
            pbar.update(len(results[str(k)][1]))

        pool_data = (
            (index, df[i : i + chunk_size], func)
            for index, i in enumerate(range(0, len(df), chunk_size))
            if index not in finished_chunks
        )


import numpy as np
import os
import pandas as pd
import requests
import shelve
import zlib

headers = {
    #'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/71.0.3578.98 Safari/537.36',
    "User-Agent": "Googlebot-Image/1.0",  # Pretend to be googlebot
    "X-Forwarded-For": "64.18.15.200",
}

def _df_split_apply(tup_arg):
    split_ind, subset, func = tup_arg
    r = subset.apply(func, axis=1)
    return (split_ind, r)

def df_multiprocess(df, processes, chunk_size, func, dataset_name):
    print("Generating parts...")
    with shelve.open(
        "%s_%s_%s_results.tmp" % (dataset_name, func.__name__, chunk_size)
    ) as results:

        pbar = tqdm(total=len(df), position=0)
        # Resume:
        finished_chunks = set([int(k) for k in results.keys()])
        pbar.desc = "Resuming"
        for k in results.keys():
            pbar.update(len(results[str(k)][1]))

        pool_data = (
            (index, df[i : i + chunk_size], func)
            for index, i in enumerate(range(0, len(df), chunk_size))
            if index not in finished_chunks
        )
        print(
            int(len(df) / chunk_size),
            "parts.",
            chunk_size,
            "per part.",
            "Using",
            processes,
            "processes",
        )

        pbar.desc = "Downloading"
        with Pool(processes) as pool:
            for i, result in enumerate(
                pool.imap_unordered(_df_split_apply, pool_data, 2)
            ):
                results[str(result[0])] = result
                pbar.update(len(result[1]))
        pbar.close()

    print("Finished Downloading.")
    return

# Unique name based on url
def _file_name(row):
    name = (
        "%s/%s_%s"
        % (
            # row["folder"],
            storage_dir,
            row.name,
            (zlib.crc32(row["url"].encode("utf-8")) & 0xFFFFFFFF),
        )
        + ".jpg"
    )
    return name

# For checking mimetypes separately without download
def check_mimetype(row):
    if os.path.isfile(str(row["file"])):


            "per part.",
            "Using",
            processes,
            "processes",
        )

        pbar.desc = "Downloading"
        with Pool(processes) as pool:
            for i, result in enumerate(
                pool.imap_unordered(_df_split_apply, pool_data, 2)
            ):
                results[str(result[0])] = result
                pbar.update(len(result[1]))
        pbar.close()

    print("Finished Downloading.")
    return

# Unique name based on url
def _file_name(row):
    name = (
        "%s/%s_%s"
        % (
            # row["folder"],
            storage_dir,
            row.name,
            (zlib.crc32(row["url"].encode("utf-8")) & 0xFFFFFFFF),
        )
        + ".jpg"
    )
    return name

# For checking mimetypes separately without download
def check_mimetype(row):
    if os.path.isfile(str(row["file"])):
        row["mimetype"] = magic.from_file(row["file"], mime=True)
        row["size"] = os.stat(row["file"]).st_size
    return row

# Don't download image, just check with a HEAD request, can't resume.
# Can use this instead of download_image to get HTTP status codes.
def check_download(row):
    fname = _file_name(row)
    try:
        # not all sites will support HEAD
        response = requests.head(
            row["url"], stream=False, timeout=5, allow_redirects=True, headers=headers
        )
        row["status"] = response.status_code


    print("Finished Downloading.")
    return

# Unique name based on url
def _file_name(row):
    name = (
        "%s/%s_%s"
        % (
            # row["folder"],
            storage_dir,
            row.name,
            (zlib.crc32(row["url"].encode("utf-8")) & 0xFFFFFFFF),
        )
        + ".jpg"
    )
    return name

# For checking mimetypes separately without download
def check_mimetype(row):
    if os.path.isfile(str(row["file"])):
        row["mimetype"] = magic.from_file(row["file"], mime=True)
        row["size"] = os.stat(row["file"]).st_size
    return row

# Don't download image, just check with a HEAD request, can't resume.
# Can use this instead of download_image to get HTTP status codes.
def check_download(row):
    fname = _file_name(row)
    try:
        # not all sites will support HEAD
        response = requests.head(
            row["url"], stream=False, timeout=5, allow_redirects=True, headers=headers
        )
        row["status"] = response.status_code
        row["headers"] = dict(response.headers)
    except:
        # log errors later, set error as 408 timeout
        row["status"] = 408
        return row
    if response.ok:
        row["file"] = fname
    return row


            # row["folder"],
            storage_dir,
            row.name,
            (zlib.crc32(row["url"].encode("utf-8")) & 0xFFFFFFFF),
        )
        + ".jpg"
    )
    return name

# For checking mimetypes separately without download
def check_mimetype(row):
    if os.path.isfile(str(row["file"])):
        row["mimetype"] = magic.from_file(row["file"], mime=True)
        row["size"] = os.stat(row["file"]).st_size
    return row

# Don't download image, just check with a HEAD request, can't resume.
# Can use this instead of download_image to get HTTP status codes.
def check_download(row):
    fname = _file_name(row)
    try:
        # not all sites will support HEAD
        response = requests.head(
            row["url"], stream=False, timeout=5, allow_redirects=True, headers=headers
        )
        row["status"] = response.status_code
        row["headers"] = dict(response.headers)
    except:
        # log errors later, set error as 408 timeout
        row["status"] = 408
        return row
    if response.ok:
        row["file"] = fname
    return row

def resize_img(req):
    image = Image.open(req).convert("RGB")
    image = TF.resize(
        # image, size=(resize_size, resize_size)
        image,
        size=resize_size,
    )  # , interpolation=Image.LANCZOS)
    return image

def download_image(row):
    fname = _file_name(row)
    # Skip Already downloaded, retry others later
    if os.path.isfile(fname):
        row["status"] = 200
        row["file"] = fname
        row["mimetype"] = magic.from_file(row["file"], mime=True)
        row["size"] = os.stat(row["file"]).st_size


# Don't download image, just check with a HEAD request, can't resume.
# Can use this instead of download_image to get HTTP status codes.
def check_download(row):
    fname = _file_name(row)
    try:
        # not all sites will support HEAD
        response = requests.head(
            row["url"], stream=False, timeout=5, allow_redirects=True, headers=headers
        )
        row["status"] = response.status_code
        row["headers"] = dict(response.headers)
    except:
        # log errors later, set error as 408 timeout
        row["status"] = 408
        return row
    if response.ok:
        row["file"] = fname
    return row

def resize_img(req):
    image = Image.open(req).convert("RGB")
    image = TF.resize(
        # image, size=(resize_size, resize_size)
        image,
        size=resize_size,
    )  # , interpolation=Image.LANCZOS)
    return image

def download_image(row):
    fname = _file_name(row)
    # Skip Already downloaded, retry others later
    if os.path.isfile(fname):
        row["status"] = 200
        row["file"] = fname
        row["mimetype"] = magic.from_file(row["file"], mime=True)
        row["size"] = os.stat(row["file"]).st_size
        return row

    try:
        # use smaller timeout to skip errors, but can result in failed downloads
        response = requests.get(
            row["url"], stream=False, timeout=5, allow_redirects=True, headers=headers
        )
        row["status"] = response.status_code
        # row['headers'] = dict(response.headers)
    except Exception as e:


        row["headers"] = dict(response.headers)
    except:
        # log errors later, set error as 408 timeout
        row["status"] = 408
        return row
    if response.ok:
        row["file"] = fname
    return row

def resize_img(req):
    image = Image.open(req).convert("RGB")
    image = TF.resize(
        # image, size=(resize_size, resize_size)
        image,
        size=resize_size,
    )  # , interpolation=Image.LANCZOS)
    return image

def download_image(row):
    fname = _file_name(row)
    # Skip Already downloaded, retry others later
    if os.path.isfile(fname):
        row["status"] = 200
        row["file"] = fname
        row["mimetype"] = magic.from_file(row["file"], mime=True)
        row["size"] = os.stat(row["file"]).st_size
        return row

    try:
        # use smaller timeout to skip errors, but can result in failed downloads
        response = requests.get(
            row["url"], stream=False, timeout=5, allow_redirects=True, headers=headers
        )
        row["status"] = response.status_code
        # row['headers'] = dict(response.headers)
    except Exception as e:
        # log errors later, set error as 408 timeout
        row["status"] = 408
        return row

    if response.ok:
        try:
            # some sites respond with gzip transport encoding
            response.raw.decode_content = True
            img = resize_img(io.BytesIO(response.content))
            img.save(fname)

            row["mimetype"] = magic.from_file(fname, mime=True)
            row["size"] = os.stat(fname).st_size

        except Exception as e:
            #     # This is if it times out during a download or decode
            row["status"] = 408

    row["file"] = fname
    return row

def open_tsv(fname, folder):
    print("Opening %s Data File..." % fname)
    df = pd.read_csv(
        fname, sep="\t", names=["url", "caption"]
    )  # , usecols=range(1, 2))
    df["folder"] = folder
    print("Processing", len(df), " Images:")
    return df

def df_from_shelve(chunk_size, func, dataset_name):
    print("Generating Dataframe from results...")
    with shelve.open(
        "%s_%s_%s_results.tmp" % (dataset_name, func.__name__, chunk_size)
    ) as results:
        keylist = sorted([int(k) for k in results.keys()])
        df = pd.concat([results[str(k)][1] for k in keylist], sort=True)
    return df


        return row

    if response.ok:
        try:
            # some sites respond with gzip transport encoding
            response.raw.decode_content = True
            img = resize_img(io.BytesIO(response.content))
            img.save(fname)

            row["mimetype"] = magic.from_file(fname, mime=True)
            row["size"] = os.stat(fname).st_size

        except Exception as e:
            #     # This is if it times out during a download or decode
            row["status"] = 408

    row["file"] = fname
    return row

def open_tsv(fname, folder):
    print("Opening %s Data File..." % fname)
    df = pd.read_csv(
        fname, sep="\t", names=["url", "caption"]
    )  # , usecols=range(1, 2))
    df["folder"] = folder
    print("Processing", len(df), " Images:")
    return df

def df_from_shelve(chunk_size, func, dataset_name):
    print("Generating Dataframe from results...")
    with shelve.open(
        "%s_%s_%s_results.tmp" % (dataset_name, func.__name__, chunk_size)
    ) as results:
        keylist = sorted([int(k) for k in results.keys()])
        df = pd.concat([results[str(k)][1] for k in keylist], sort=True)
    return df

resize_size = 384

config_path = get_abs_path("configs/datasets/conceptual_caption/defaults_12m.yaml")

storage_dir = OmegaConf.load(
    config_path
).datasets.conceptual_caption_12m.build_info.images.storage
storage_dir = Path(get_cache_path(storage_dir))


            row["size"] = os.stat(fname).st_size

        except Exception as e:
            #     # This is if it times out during a download or decode
            row["status"] = 408

    row["file"] = fname
    return row

def open_tsv(fname, folder):
    print("Opening %s Data File..." % fname)
    df = pd.read_csv(
        fname, sep="\t", names=["url", "caption"]
    )  # , usecols=range(1, 2))
    df["folder"] = folder
    print("Processing", len(df), " Images:")
    return df

def df_from_shelve(chunk_size, func, dataset_name):
    print("Generating Dataframe from results...")
    with shelve.open(
        "%s_%s_%s_results.tmp" % (dataset_name, func.__name__, chunk_size)
    ) as results:
        keylist = sorted([int(k) for k in results.keys()])
        df = pd.concat([results[str(k)][1] for k in keylist], sort=True)
    return df

resize_size = 384

config_path = get_abs_path("configs/datasets/conceptual_caption/defaults_12m.yaml")

storage_dir = OmegaConf.load(
    config_path
).datasets.conceptual_caption_12m.build_info.images.storage
storage_dir = Path(get_cache_path(storage_dir))

os.makedirs(storage_dir, exist_ok=True)

# number of processes in the pool can be larger than cores
num_processes = 96
# num_processes = 1
# chunk_size is how many images per chunk per process - changing this resets progress when restarting.
images_per_part = 100

data_name = "cc12m"


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import webdataset as wds
from lavis.datasets.datasets.base_dataset import BaseDataset

class LaionDataset(BaseDataset):
    def __init__(self, vis_processor, text_processor, location):
        super().__init__(vis_processor=vis_processor, text_processor=text_processor)

        self.inner_dataset = wds.DataPipeline(
            wds.ResampledShards(location),
            wds.tarfile_to_samples(handler=wds.warn_and_continue),
            wds.shuffle(1000, handler=wds.warn_and_continue),
            wds.decode("pilrgb", handler=wds.warn_and_continue),
            wds.to_tuple("jpg", "json", handler=wds.warn_and_continue),
            wds.map_tuple(self.vis_processor, handler=wds.warn_and_continue),
            wds.map(self.to_dict, handler=wds.warn_and_continue),
        )

    def to_dict(self, sample):
        return {
            "image": sample[0],
            "text_input": self.text_processor(sample[1]["caption"]),
        }

if __name__ == "__main__":
    from torchvision import transforms

    def to_image_text_pair(sample):
        return sample[0], sample[1]["caption"]

    normalize = transforms.Normalize(
        (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
    )

    transform_train = transforms.Compose(
        [
            transforms.RandomResizedCrop(256, scale=(0.2, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
    )

    dataset = LaionDataset(
        vis_processor=transform_train,
        text_processor=lambda x: x,
        location="/export/laion/laion2B-multi/part-00000/{00000..01743}.tar",
    )

    import torch

    loader = torch.utils.data.DataLoader(dataset.inner_dataset, batch_size=2)

    print(next(iter(loader))["text_input"])


LaionDataset

"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import webdataset as wds
from lavis.datasets.datasets.base_dataset import BaseDataset

class LaionDataset(BaseDataset):
    def __init__(self, vis_processor, text_processor, location):
        super().__init__(vis_processor=vis_processor, text_processor=text_processor)

        self.inner_dataset = wds.DataPipeline(
            wds.ResampledShards(location),
            wds.tarfile_to_samples(handler=wds.warn_and_continue),
            wds.shuffle(1000, handler=wds.warn_and_continue),
            wds.decode("pilrgb", handler=wds.warn_and_continue),
            wds.to_tuple("jpg", "json", handler=wds.warn_and_continue),
            wds.map_tuple(self.vis_processor, handler=wds.warn_and_continue),
            wds.map(self.to_dict, handler=wds.warn_and_continue),
        )

    def to_dict(self, sample):
        return {
            "image": sample[0],
            "text_input": self.text_processor(sample[1]["caption"]),
        }

if __name__ == "__main__":
    from torchvision import transforms

    def to_image_text_pair(sample):
        return sample[0], sample[1]["caption"]

    normalize = transforms.Normalize(
        (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
    )

    transform_train = transforms.Compose(
        [
            transforms.RandomResizedCrop(256, scale=(0.2, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
    )


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import webdataset as wds
from lavis.datasets.datasets.base_dataset import BaseDataset

class LaionDataset(BaseDataset):
    def __init__(self, vis_processor, text_processor, location):
        super().__init__(vis_processor=vis_processor, text_processor=text_processor)

        self.inner_dataset = wds.DataPipeline(
            wds.ResampledShards(location),
            wds.tarfile_to_samples(handler=wds.warn_and_continue),
            wds.shuffle(1000, handler=wds.warn_and_continue),
            wds.decode("pilrgb", handler=wds.warn_and_continue),
            wds.to_tuple("jpg", "json", handler=wds.warn_and_continue),
            wds.map_tuple(self.vis_processor, handler=wds.warn_and_continue),
            wds.map(self.to_dict, handler=wds.warn_and_continue),
        )

    def to_dict(self, sample):
        return {
            "image": sample[0],
            "text_input": self.text_processor(sample[1]["caption"]),
        }

if __name__ == "__main__":
    from torchvision import transforms

    def to_image_text_pair(sample):
        return sample[0], sample[1]["caption"]

    normalize = transforms.Normalize(
        (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
    )

    transform_train = transforms.Compose(
        [


"""

import webdataset as wds
from lavis.datasets.datasets.base_dataset import BaseDataset

class LaionDataset(BaseDataset):
    def __init__(self, vis_processor, text_processor, location):
        super().__init__(vis_processor=vis_processor, text_processor=text_processor)

        self.inner_dataset = wds.DataPipeline(
            wds.ResampledShards(location),
            wds.tarfile_to_samples(handler=wds.warn_and_continue),
            wds.shuffle(1000, handler=wds.warn_and_continue),
            wds.decode("pilrgb", handler=wds.warn_and_continue),
            wds.to_tuple("jpg", "json", handler=wds.warn_and_continue),
            wds.map_tuple(self.vis_processor, handler=wds.warn_and_continue),
            wds.map(self.to_dict, handler=wds.warn_and_continue),
        )

    def to_dict(self, sample):
        return {
            "image": sample[0],
            "text_input": self.text_processor(sample[1]["caption"]),
        }

if __name__ == "__main__":
    from torchvision import transforms

    def to_image_text_pair(sample):
        return sample[0], sample[1]["caption"]

    normalize = transforms.Normalize(
        (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
    )

    transform_train = transforms.Compose(
        [
            transforms.RandomResizedCrop(256, scale=(0.2, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
    )


        self.inner_dataset = wds.DataPipeline(
            wds.ResampledShards(location),
            wds.tarfile_to_samples(handler=wds.warn_and_continue),
            wds.shuffle(1000, handler=wds.warn_and_continue),
            wds.decode("pilrgb", handler=wds.warn_and_continue),
            wds.to_tuple("jpg", "json", handler=wds.warn_and_continue),
            wds.map_tuple(self.vis_processor, handler=wds.warn_and_continue),
            wds.map(self.to_dict, handler=wds.warn_and_continue),
        )

    def to_dict(self, sample):
        return {
            "image": sample[0],
            "text_input": self.text_processor(sample[1]["caption"]),
        }

if __name__ == "__main__":
    from torchvision import transforms

    def to_image_text_pair(sample):
        return sample[0], sample[1]["caption"]

    normalize = transforms.Normalize(
        (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
    )

    transform_train = transforms.Compose(
        [
            transforms.RandomResizedCrop(256, scale=(0.2, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
    )

    dataset = LaionDataset(
        vis_processor=transform_train,
        text_processor=lambda x: x,
        location="/export/laion/laion2B-multi/part-00000/{00000..01743}.tar",
    )


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import os
from collections import OrderedDict

from lavis.datasets.datasets.base_dataset import BaseDataset
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]
        visual_key = "image" if "image" in ann else "video"

        return OrderedDict(
            {
                "file": ann[visual_key],
                "caption": ann["caption"],
                visual_key: sample[visual_key],
            }
        )

class RetrievalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {
            "image": image,
            "text_input": caption,
            "image_id": self.img_ids[ann["image_id"]],
            "instance_id": ann["instance_id"],
        }

class RetrievalEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.text = []
        self.image = []
        self.txt2img = {}
        self.img2txt = {}

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["image"])
            self.img2txt[img_id] = []
            for i, caption in enumerate(ann["caption"]):
                self.text.append(self.text_processor(caption))
                self.img2txt[img_id].append(txt_id)
                self.txt2img[txt_id] = img_id
                txt_id += 1

    def __getitem__(self, index):

        image_path = os.path.join(self.vis_root, self.annotation[index]["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)

        return {"image": image, "index": index}

class VideoRetrievalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of videos.
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["video"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        vpath = os.path.join(self.vis_root, ann["video"])

        video = self.vis_processor(vpath)
        caption = self.text_processor(ann["caption"])

        # return image, caption, self.img_ids[ann['image_id']]
        return {
            "video": video,
            "text_input": caption,
            "image_id": self.img_ids[ann["video"]],
        }

class VideoRetrievalEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of videos.
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.text = []
        self.image = []
        self.txt2img = {}
        self.img2txt = {}

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["video"])
            self.img2txt[img_id] = []
            for i, caption in enumerate(ann["caption"]):
                self.text.append(self.text_processor(caption))
                self.img2txt[img_id].append(txt_id)
                self.txt2img[txt_id] = img_id
                txt_id += 1

    def __getitem__(self, index):
        ann = self.annotation[index]

        vpath = os.path.join(self.vis_root, ann["video"])
        video = self.vis_processor(vpath)

        return {"video": video, "index": index}


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import os
from collections import OrderedDict

from lavis.datasets.datasets.base_dataset import BaseDataset
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]
        visual_key = "image" if "image" in ann else "video"

        return OrderedDict(
            {
                "file": ann[visual_key],
                "caption": ann["caption"],
                visual_key: sample[visual_key],
            }
        )

class RetrievalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):


RetrievalDataset

from collections import OrderedDict

from lavis.datasets.datasets.base_dataset import BaseDataset
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]
        visual_key = "image" if "image" in ann else "video"

        return OrderedDict(
            {
                "file": ann[visual_key],
                "caption": ann["caption"],
                visual_key: sample[visual_key],
            }
        )

class RetrievalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {
            "image": image,
            "text_input": caption,
            "image_id": self.img_ids[ann["image_id"]],
            "instance_id": ann["instance_id"],
        }

class RetrievalEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.text = []
        self.image = []
        self.txt2img = {}
        self.img2txt = {}

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["image"])


RetrievalEvalDataset

                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {
            "image": image,
            "text_input": caption,
            "image_id": self.img_ids[ann["image_id"]],
            "instance_id": ann["instance_id"],
        }

class RetrievalEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.text = []
        self.image = []
        self.txt2img = {}
        self.img2txt = {}

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["image"])
            self.img2txt[img_id] = []
            for i, caption in enumerate(ann["caption"]):
                self.text.append(self.text_processor(caption))
                self.img2txt[img_id].append(txt_id)
                self.txt2img[txt_id] = img_id
                txt_id += 1

    def __getitem__(self, index):

        image_path = os.path.join(self.vis_root, self.annotation[index]["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)

        return {"image": image, "index": index}

class VideoRetrievalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of videos.
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["video"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):


VideoRetrievalDataset

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["image"])
            self.img2txt[img_id] = []
            for i, caption in enumerate(ann["caption"]):
                self.text.append(self.text_processor(caption))
                self.img2txt[img_id].append(txt_id)
                self.txt2img[txt_id] = img_id
                txt_id += 1

    def __getitem__(self, index):

        image_path = os.path.join(self.vis_root, self.annotation[index]["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)

        return {"image": image, "index": index}

class VideoRetrievalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of videos.
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["video"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        vpath = os.path.join(self.vis_root, ann["video"])

        video = self.vis_processor(vpath)
        caption = self.text_processor(ann["caption"])

        # return image, caption, self.img_ids[ann['image_id']]
        return {
            "video": video,
            "text_input": caption,
            "image_id": self.img_ids[ann["video"]],
        }

class VideoRetrievalEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of videos.
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.text = []
        self.image = []
        self.txt2img = {}
        self.img2txt = {}

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["video"])


VideoRetrievalEvalDataset

                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        vpath = os.path.join(self.vis_root, ann["video"])

        video = self.vis_processor(vpath)
        caption = self.text_processor(ann["caption"])

        # return image, caption, self.img_ids[ann['image_id']]
        return {
            "video": video,
            "text_input": caption,
            "image_id": self.img_ids[ann["video"]],
        }

class VideoRetrievalEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of videos.
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.text = []
        self.image = []
        self.txt2img = {}
        self.img2txt = {}

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["video"])
            self.img2txt[img_id] = []
            for i, caption in enumerate(ann["caption"]):
                self.text.append(self.text_processor(caption))
                self.img2txt[img_id].append(txt_id)
                self.txt2img[txt_id] = img_id
                txt_id += 1

    def __getitem__(self, index):
        ann = self.annotation[index]

        vpath = os.path.join(self.vis_root, ann["video"])
        video = self.vis_processor(vpath)

        return {"video": video, "index": index}


                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        vpath = os.path.join(self.vis_root, ann["video"])

        video = self.vis_processor(vpath)
        caption = self.text_processor(ann["caption"])

        # return image, caption, self.img_ids[ann['image_id']]
        return {
            "video": video,
            "text_input": caption,
            "image_id": self.img_ids[ann["video"]],
        }

class VideoRetrievalEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of videos.
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.text = []
        self.image = []
        self.txt2img = {}
        self.img2txt = {}

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["video"])
            self.img2txt[img_id] = []
            for i, caption in enumerate(ann["caption"]):
                self.text.append(self.text_processor(caption))
                self.img2txt[img_id].append(txt_id)
                self.txt2img[txt_id] = img_id
                txt_id += 1

    def __getitem__(self, index):
        ann = self.annotation[index]

        vpath = os.path.join(self.vis_root, ann["video"])
        video = self.vis_processor(vpath)

        return {"video": video, "index": index}


        split (string): val or test
        """

        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.text = []
        self.image = []
        self.txt2img = {}
        self.img2txt = {}

        txt_id = 0
        for img_id, ann in enumerate(self.annotation):
            self.image.append(ann["video"])
            self.img2txt[img_id] = []
            for i, caption in enumerate(ann["caption"]):
                self.text.append(self.text_processor(caption))
                self.img2txt[img_id].append(txt_id)
                self.txt2img[txt_id] = img_id
                txt_id += 1

    def __getitem__(self, index):
        ann = self.annotation[index]

        vpath = os.path.join(self.vis_root, ann["video"])
        video = self.vis_processor(vpath)

        return {"video": video, "index": index}


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import os
from collections import OrderedDict

from lavis.datasets.datasets.base_dataset import BaseDataset
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file": os.path.basename(ann["image"]),
                "caption": ann["caption"],
                "image": sample["image"],
            }
        )

class ImageTextPairDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

    def __getitem__(self, index):

        # TODO this assumes image input, not general enough
        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {"image": image, "text_input": caption}


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import os
from collections import OrderedDict

from lavis.datasets.datasets.base_dataset import BaseDataset
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file": os.path.basename(ann["image"]),
                "caption": ann["caption"],
                "image": sample["image"],
            }
        )

class ImageTextPairDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

    def __getitem__(self, index):

        # TODO this assumes image input, not general enough
        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])


ImageTextPairDataset

import os
from collections import OrderedDict

from lavis.datasets.datasets.base_dataset import BaseDataset
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file": os.path.basename(ann["image"]),
                "caption": ann["caption"],
                "image": sample["image"],
            }
        )

class ImageTextPairDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

    def __getitem__(self, index):

        # TODO this assumes image input, not general enough
        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {"image": image, "text_input": caption}


from collections import OrderedDict

from lavis.datasets.datasets.base_dataset import BaseDataset
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file": os.path.basename(ann["image"]),
                "caption": ann["caption"],
                "image": sample["image"],
            }
        )

class ImageTextPairDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

    def __getitem__(self, index):

        # TODO this assumes image input, not general enough
        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {"image": image, "text_input": caption}


    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file": os.path.basename(ann["image"]),
                "caption": ann["caption"],
                "image": sample["image"],
            }
        )

class ImageTextPairDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

    def __getitem__(self, index):

        # TODO this assumes image input, not general enough
        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {"image": image, "text_input": caption}


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import torch
import copy
import os
import random
import json
import ast
import re
from PIL import Image
from lavis.datasets.datasets.base_dataset import BaseDataset

class MusicAVQADataset(BaseDataset):
    def __init__(self, **kwargs):
        super().__init__(kwargs['vis_processor'], kwargs['text_processor'], kwargs['vis_root'], kwargs['ann_paths'], kwargs['modalities'])

        self.modalities = kwargs['modalities']

        for modality in self.modalities:
            if 'image' in modality:
                setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())
                continue
            setattr(self, f"{modality}_root", kwargs[f"{modality}_root"])
            setattr(self, f"{modality}_processor", kwargs[f"{modality}_processor"])
            setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())

        self.sample_ids = set.intersection(*[set(getattr(self, f"existing_{modality}_annotation")) for modality in self.modalities])

        try:
            self.annotation = [ann for ann in self.annotation if ann['video_id'] in self.sample_ids]
        except:
            self.sample_ids = set.intersection(*[set(self.get_existing_audio_pt_annotations()) for modality in self.modalities])
            self.annotation = [ann for ann in self.annotation if ann['id'] in self.sample_ids]

    def get_existing_audio_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.audio_root)]

    def get_existing_frame_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.frame_root)]

    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')

    def get_video_path(self, ann):
        return os.path.join(self.video_root, f'{ann["video_id"]}.mp4')

    def get_flow_path(self, ann):
        return os.path.join(self.flow_root, f'{ann["video_id"]}/')

    def get_frame_path(self, ann):
        return os.path.join(self.frame_root, f'{ann["video_id"]}/')

    def __getitem__(self, index):
        ann = copy.deepcopy(self.annotation[index])

        for modality in self.modalities:
            if 'id' not in ann:
                ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)
                if type(ann[f"{modality}_path"]) == list:
                    ann[f"{modality}_path"] = random.choice(ann[f"{modality}_path"])

            if 'image' in modality:
                ann['rgb'] = self.vis_processor(Image.open(ann[f"images_path"]))
            else:
                if modality == 'video':
                    rgb, indices, fps = getattr(self, f"{modality}_processor")(ann[f"{modality}_path"])
                    ann['rgb'] = rgb.to(torch.float32)

                if modality == 'frame':
                    indices, clip = None, None
                    ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)
                    frms, indices = getattr(self, f"{modality}_processor")(ann[f"{modality}_path"], clip_proposal=clip, indices=indices)
                    rgb = frms.permute(1, 0, 2, 3)
                    assert len(rgb) == getattr(self, f"{modality}_processor").n_frms
                    ann['rgb'] = rgb

                if modality == 'flow':
                    assert indices is not None
                    ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)
                    flow, _ = getattr(self, f"{modality}_processor")(ann[f"{modality}_path"], clip_proposal=clip, indices=indices, type='depth')
                    ann['flow'] = flow

                if modality == 'audio':
                    if 'id' in ann: # aduio pt data
                        ann[f"{modality}_path"] = self.get_pt_audio_path(ann)
                    ann[modality] = getattr(self, f"{modality}_processor")(ann[f"{modality}_path"]).to(torch.float32)

        if 'id' not in ann:
            ann["sample_id"] = ann["video_id"]
            if len(ann['templ_values']) != 0:
                question = ann['question_content']
                templ_values = ast.literal_eval(ann['templ_values'])
                matches = re.findall(r'<(.*?)>', question)
                for k, v in zip(matches, templ_values):
                    question = question.replace('<'+k+'>', v)
            else:
                question = ann['question_content']
            ann['qa_input'] =  'Question: ' + self.text_processor(question) + ' Based on the frames and audio information, answer the question using a single word or phase.'
            ann["question_id"] = ann['question_id']
            answers = ann['anser']
            if '_' in answers:
                answers = answers.replace('_', ' ')
            ann['answers'] = answers #ann['anser']

        else:
            ann["sample_id"] = ann["id"]
            ann['answers'] = ann['answer']
            ann["question_id"] = ann['id']
            question = ann['question']
            ann['qa_input'] =  'Question: ' + self.text_processor(question) + ' Based on the audio information, answer the question using a single word or phase.'

        return ann

class MusicAVQAInstructDataset(MusicAVQADataset):
    def __getitem__(self, index):
        data = super().__getitem__(index)
        if data != None:
            data['answer'] = data["answers"] # needed to use gqa task
            data['qa_output'] = data["answers"]

        return data


MusicAVQAInstructDataset

                for k, v in zip(matches, templ_values):
                    question = question.replace('<'+k+'>', v)
            else:
                question = ann['question_content']
            ann['qa_input'] =  'Question: ' + self.text_processor(question) + ' Based on the frames and audio information, answer the question using a single word or phase.'
            ann["question_id"] = ann['question_id']
            answers = ann['anser']
            if '_' in answers:
                answers = answers.replace('_', ' ')
            ann['answers'] = answers #ann['anser']

        else:
            ann["sample_id"] = ann["id"]
            ann['answers'] = ann['answer']
            ann["question_id"] = ann['id']
            question = ann['question']
            ann['qa_input'] =  'Question: ' + self.text_processor(question) + ' Based on the audio information, answer the question using a single word or phase.'

        return ann

class MusicAVQAInstructDataset(MusicAVQADataset):
    def __getitem__(self, index):
        data = super().__getitem__(index)
        if data != None:
            data['answer'] = data["answers"] # needed to use gqa task
            data['qa_output'] = data["answers"]

        return data


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import torch
import copy
import os
import random
import json
import ast
import re
from PIL import Image
from lavis.datasets.datasets.base_dataset import BaseDataset

class MusicAVQADataset(BaseDataset):
    def __init__(self, **kwargs):
        super().__init__(kwargs['vis_processor'], kwargs['text_processor'], kwargs['vis_root'], kwargs['ann_paths'], kwargs['modalities'])

        self.modalities = kwargs['modalities']

        for modality in self.modalities:
            if 'image' in modality:
                setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())
                continue
            setattr(self, f"{modality}_root", kwargs[f"{modality}_root"])
            setattr(self, f"{modality}_processor", kwargs[f"{modality}_processor"])
            setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())

        self.sample_ids = set.intersection(*[set(getattr(self, f"existing_{modality}_annotation")) for modality in self.modalities])

        try:
            self.annotation = [ann for ann in self.annotation if ann['video_id'] in self.sample_ids]
        except:
            self.sample_ids = set.intersection(*[set(self.get_existing_audio_pt_annotations()) for modality in self.modalities])
            self.annotation = [ann for ann in self.annotation if ann['id'] in self.sample_ids]

    def get_existing_audio_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.audio_root)]

    def get_existing_frame_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.frame_root)]

    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):


        super().__init__(kwargs['vis_processor'], kwargs['text_processor'], kwargs['vis_root'], kwargs['ann_paths'], kwargs['modalities'])

        self.modalities = kwargs['modalities']

        for modality in self.modalities:
            if 'image' in modality:
                setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())
                continue
            setattr(self, f"{modality}_root", kwargs[f"{modality}_root"])
            setattr(self, f"{modality}_processor", kwargs[f"{modality}_processor"])
            setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())

        self.sample_ids = set.intersection(*[set(getattr(self, f"existing_{modality}_annotation")) for modality in self.modalities])

        try:
            self.annotation = [ann for ann in self.annotation if ann['video_id'] in self.sample_ids]
        except:
            self.sample_ids = set.intersection(*[set(self.get_existing_audio_pt_annotations()) for modality in self.modalities])
            self.annotation = [ann for ann in self.annotation if ann['id'] in self.sample_ids]

    def get_existing_audio_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.audio_root)]

    def get_existing_frame_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.frame_root)]

    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')


        for modality in self.modalities:
            if 'image' in modality:
                setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())
                continue
            setattr(self, f"{modality}_root", kwargs[f"{modality}_root"])
            setattr(self, f"{modality}_processor", kwargs[f"{modality}_processor"])
            setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())

        self.sample_ids = set.intersection(*[set(getattr(self, f"existing_{modality}_annotation")) for modality in self.modalities])

        try:
            self.annotation = [ann for ann in self.annotation if ann['video_id'] in self.sample_ids]
        except:
            self.sample_ids = set.intersection(*[set(self.get_existing_audio_pt_annotations()) for modality in self.modalities])
            self.annotation = [ann for ann in self.annotation if ann['id'] in self.sample_ids]

    def get_existing_audio_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.audio_root)]

    def get_existing_frame_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.frame_root)]

    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')

    def get_video_path(self, ann):
        return os.path.join(self.video_root, f'{ann["video_id"]}.mp4')


                setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())
                continue
            setattr(self, f"{modality}_root", kwargs[f"{modality}_root"])
            setattr(self, f"{modality}_processor", kwargs[f"{modality}_processor"])
            setattr(self, f"existing_{modality}_annotation",getattr(self, f'get_existing_{modality}_annotations')())

        self.sample_ids = set.intersection(*[set(getattr(self, f"existing_{modality}_annotation")) for modality in self.modalities])

        try:
            self.annotation = [ann for ann in self.annotation if ann['video_id'] in self.sample_ids]
        except:
            self.sample_ids = set.intersection(*[set(self.get_existing_audio_pt_annotations()) for modality in self.modalities])
            self.annotation = [ann for ann in self.annotation if ann['id'] in self.sample_ids]

    def get_existing_audio_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.audio_root)]

    def get_existing_frame_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.frame_root)]

    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')

    def get_video_path(self, ann):
        return os.path.join(self.video_root, f'{ann["video_id"]}.mp4')

    def get_flow_path(self, ann):
        return os.path.join(self.flow_root, f'{ann["video_id"]}/')


            self.annotation = [ann for ann in self.annotation if ann['video_id'] in self.sample_ids]
        except:
            self.sample_ids = set.intersection(*[set(self.get_existing_audio_pt_annotations()) for modality in self.modalities])
            self.annotation = [ann for ann in self.annotation if ann['id'] in self.sample_ids]

    def get_existing_audio_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.audio_root)]

    def get_existing_frame_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.frame_root)]

    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')

    def get_video_path(self, ann):
        return os.path.join(self.video_root, f'{ann["video_id"]}.mp4')

    def get_flow_path(self, ann):
        return os.path.join(self.flow_root, f'{ann["video_id"]}/')

    def get_frame_path(self, ann):
        return os.path.join(self.frame_root, f'{ann["video_id"]}/')

    def __getitem__(self, index):
        ann = copy.deepcopy(self.annotation[index])

        for modality in self.modalities:
            if 'id' not in ann:
                ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)


            self.annotation = [ann for ann in self.annotation if ann['id'] in self.sample_ids]

    def get_existing_audio_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.audio_root)]

    def get_existing_frame_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.frame_root)]

    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')

    def get_video_path(self, ann):
        return os.path.join(self.video_root, f'{ann["video_id"]}.mp4')

    def get_flow_path(self, ann):
        return os.path.join(self.flow_root, f'{ann["video_id"]}/')

    def get_frame_path(self, ann):
        return os.path.join(self.frame_root, f'{ann["video_id"]}/')

    def __getitem__(self, index):
        ann = copy.deepcopy(self.annotation[index])

        for modality in self.modalities:
            if 'id' not in ann:
                ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)
                if type(ann[f"{modality}_path"]) == list:
                    ann[f"{modality}_path"] = random.choice(ann[f"{modality}_path"])

            if 'image' in modality:


    def get_existing_frame_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.frame_root)]

    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')

    def get_video_path(self, ann):
        return os.path.join(self.video_root, f'{ann["video_id"]}.mp4')

    def get_flow_path(self, ann):
        return os.path.join(self.flow_root, f'{ann["video_id"]}/')

    def get_frame_path(self, ann):
        return os.path.join(self.frame_root, f'{ann["video_id"]}/')

    def __getitem__(self, index):
        ann = copy.deepcopy(self.annotation[index])

        for modality in self.modalities:
            if 'id' not in ann:
                ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)
                if type(ann[f"{modality}_path"]) == list:
                    ann[f"{modality}_path"] = random.choice(ann[f"{modality}_path"])

            if 'image' in modality:
                ann['rgb'] = self.vis_processor(Image.open(ann[f"images_path"]))
            else:
                if modality == 'video':


    def get_existing_flow_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.flow_root)]

    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')

    def get_video_path(self, ann):
        return os.path.join(self.video_root, f'{ann["video_id"]}.mp4')

    def get_flow_path(self, ann):
        return os.path.join(self.flow_root, f'{ann["video_id"]}/')

    def get_frame_path(self, ann):
        return os.path.join(self.frame_root, f'{ann["video_id"]}/')

    def __getitem__(self, index):
        ann = copy.deepcopy(self.annotation[index])

        for modality in self.modalities:
            if 'id' not in ann:
                ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)
                if type(ann[f"{modality}_path"]) == list:
                    ann[f"{modality}_path"] = random.choice(ann[f"{modality}_path"])

            if 'image' in modality:
                ann['rgb'] = self.vis_processor(Image.open(ann[f"images_path"]))
            else:
                if modality == 'video':
                    rgb, indices, fps = getattr(self, f"{modality}_processor")(ann[f"{modality}_path"])
                    ann['rgb'] = rgb.to(torch.float32)


    # def get_existing_audio_pt_annotations(self):
    #     return [f.split('.')[0] for f in os.listdir('/nas-ssd2/shoubin/datasets/audios')]

    # def get_pt_audio_path(self, ann):
    #     return os.path.join('/nas-ssd2/shoubin/datasets/audios', f'{str(ann["id"])}.wav')

    def get_existing_video_annotations(self):
        return [f.split('.')[0] for f in os.listdir(self.video_root)]

    def get_audio_path(self, ann):
        # return os.path.join(self.audio_root, f'{ann["video_id"]}.flac')
        return os.path.join(self.audio_root, f'{ann["video_id"]}.mp4')

    def get_video_path(self, ann):
        return os.path.join(self.video_root, f'{ann["video_id"]}.mp4')

    def get_flow_path(self, ann):
        return os.path.join(self.flow_root, f'{ann["video_id"]}/')

    def get_frame_path(self, ann):
        return os.path.join(self.frame_root, f'{ann["video_id"]}/')

    def __getitem__(self, index):
        ann = copy.deepcopy(self.annotation[index])

        for modality in self.modalities:
            if 'id' not in ann:
                ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)
                if type(ann[f"{modality}_path"]) == list:
                    ann[f"{modality}_path"] = random.choice(ann[f"{modality}_path"])

            if 'image' in modality:
                ann['rgb'] = self.vis_processor(Image.open(ann[f"images_path"]))
            else:
                if modality == 'video':
                    rgb, indices, fps = getattr(self, f"{modality}_processor")(ann[f"{modality}_path"])
                    ann['rgb'] = rgb.to(torch.float32)

                if modality == 'frame':
                    indices, clip = None, None
                    ann[f"{modality}_path"] = getattr(self, f"get_{modality}_path")(ann)


                    question = question.replace('<'+k+'>', v)
            else:
                question = ann['question_content']
            ann['qa_input'] =  'Question: ' + self.text_processor(question) + ' Based on the frames and audio information, answer the question using a single word or phase.'
            ann["question_id"] = ann['question_id']
            answers = ann['anser']
            if '_' in answers:
                answers = answers.replace('_', ' ')
            ann['answers'] = answers #ann['anser']

        else:
            ann["sample_id"] = ann["id"]
            ann['answers'] = ann['answer']
            ann["question_id"] = ann['id']
            question = ann['question']
            ann['qa_input'] =  'Question: ' + self.text_processor(question) + ' Based on the audio information, answer the question using a single word or phase.'

        return ann

class MusicAVQAInstructDataset(MusicAVQADataset):
    def __getitem__(self, index):
        data = super().__getitem__(index)
        if data != None:
            data['answer'] = data["answers"] # needed to use gqa task
            data['qa_output'] = data["answers"]

        return data


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import os
from collections import OrderedDict

from PIL import Image

from lavis.datasets.datasets.base_dataset import BaseDataset

import json
import copy

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file": ann["image"],
                "dialogue": ann["dialogue"],
                "image": sample["image"],
            }
        )

class DialogueDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """

        self.vis_root = vis_root

        self.annotation = []
        for ann_path in ann_paths:
            dialogs = json.load(open(ann_path, "r"))["dialogs"]
            for dialog in dialogs:
                all_turns = dialog["dialog"]
                dialogue_context = []
                for turn in all_turns:
                    dialog_instance = copy.deepcopy(dialog)
                    question = turn["question"]
                    answer = turn["answer"]

                    dialog_instance["dialog"] = copy.deepcopy(dialogue_context)
                    dialog_instance["question"] = question
                    dialog_instance["answer"] = answer
                    self.annotation.append(dialog_instance)
                    dialogue_context.append(turn)

        self.vis_processor = vis_processor
        self.text_processor = text_processor

        self._add_instance_ids()

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {
            "image": image,
            "text_input": caption,
            "image_id": self.img_ids[ann["image_id"]],
        }

class DialogueEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        self.vis_root = vis_root

        self.annotation = []
        for ann_path in ann_paths:
            dialogs = json.load(open(ann_path, "r"))["dialogs"]
            for dialog in dialogs:
                all_turns = dialog["dialog"]
                dialogue_context = all_turns[:-1]
                last_turn = all_turns[-1]

                question = last_turn["question"]
                answer = last_turn["answer"]

                dialog["dialog"] = dialogue_context
                dialog["question"] = question
                dialog["answer"] = answer

                self.annotation.append(dialog)

        self.vis_processor = vis_processor
        self.text_processor = text_processor

        self._add_instance_ids()

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)

        return {
            "image": image,
            "image_id": ann["image_id"],
            "instance_id": ann["instance_id"],
        }


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import os
from collections import OrderedDict

from PIL import Image

from lavis.datasets.datasets.base_dataset import BaseDataset

import json
import copy

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file": ann["image"],
                "dialogue": ann["dialogue"],
                "image": sample["image"],
            }
        )

class DialogueDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """

        self.vis_root = vis_root

        self.annotation = []
        for ann_path in ann_paths:
            dialogs = json.load(open(ann_path, "r"))["dialogs"]
            for dialog in dialogs:
                all_turns = dialog["dialog"]
                dialogue_context = []
                for turn in all_turns:
                    dialog_instance = copy.deepcopy(dialog)
                    question = turn["question"]


DialogueDataset

from lavis.datasets.datasets.base_dataset import BaseDataset

import json
import copy

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file": ann["image"],
                "dialogue": ann["dialogue"],
                "image": sample["image"],
            }
        )

class DialogueDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        """

        self.vis_root = vis_root

        self.annotation = []
        for ann_path in ann_paths:
            dialogs = json.load(open(ann_path, "r"))["dialogs"]
            for dialog in dialogs:
                all_turns = dialog["dialog"]
                dialogue_context = []
                for turn in all_turns:
                    dialog_instance = copy.deepcopy(dialog)
                    question = turn["question"]
                    answer = turn["answer"]

                    dialog_instance["dialog"] = copy.deepcopy(dialogue_context)
                    dialog_instance["question"] = question
                    dialog_instance["answer"] = answer
                    self.annotation.append(dialog_instance)
                    dialogue_context.append(turn)

        self.vis_processor = vis_processor
        self.text_processor = text_processor

        self._add_instance_ids()

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {
            "image": image,
            "text_input": caption,
            "image_id": self.img_ids[ann["image_id"]],
        }

class DialogueEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        self.vis_root = vis_root

        self.annotation = []
        for ann_path in ann_paths:
            dialogs = json.load(open(ann_path, "r"))["dialogs"]
            for dialog in dialogs:
                all_turns = dialog["dialog"]
                dialogue_context = all_turns[:-1]
                last_turn = all_turns[-1]


DialogueEvalDataset

                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {
            "image": image,
            "text_input": caption,
            "image_id": self.img_ids[ann["image_id"]],
        }

class DialogueEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        self.vis_root = vis_root

        self.annotation = []
        for ann_path in ann_paths:
            dialogs = json.load(open(ann_path, "r"))["dialogs"]
            for dialog in dialogs:
                all_turns = dialog["dialog"]
                dialogue_context = all_turns[:-1]
                last_turn = all_turns[-1]

                question = last_turn["question"]
                answer = last_turn["answer"]

                dialog["dialog"] = dialogue_context
                dialog["question"] = question
                dialog["answer"] = answer

                self.annotation.append(dialog)

        self.vis_processor = vis_processor
        self.text_processor = text_processor

        self._add_instance_ids()

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)

        return {
            "image": image,
            "image_id": ann["image_id"],
            "instance_id": ann["instance_id"],
        }


                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {
            "image": image,
            "text_input": caption,
            "image_id": self.img_ids[ann["image_id"]],
        }

class DialogueEvalDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        """
        vis_root (string): Root directory of images (e.g. coco/images/)
        ann_root (string): directory to store the annotation file
        split (string): val or test
        """

        self.vis_root = vis_root

        self.annotation = []
        for ann_path in ann_paths:
            dialogs = json.load(open(ann_path, "r"))["dialogs"]
            for dialog in dialogs:
                all_turns = dialog["dialog"]
                dialogue_context = all_turns[:-1]
                last_turn = all_turns[-1]

                question = last_turn["question"]
                answer = last_turn["answer"]

                dialog["dialog"] = dialogue_context
                dialog["question"] = question
                dialog["answer"] = answer

                self.annotation.append(dialog)

        self.vis_processor = vis_processor
        self.text_processor = text_processor

        self._add_instance_ids()

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)

        return {
            "image": image,
            "image_id": ann["image_id"],
            "instance_id": ann["instance_id"],
        }


                dialog["dialog"] = dialogue_context
                dialog["question"] = question
                dialog["answer"] = answer

                self.annotation.append(dialog)

        self.vis_processor = vis_processor
        self.text_processor = text_processor

        self._add_instance_ids()

        self.img_ids = {}
        n = 0
        for ann in self.annotation:
            img_id = ann["image_id"]
            if img_id not in self.img_ids.keys():
                self.img_ids[img_id] = n
                n += 1

    def __getitem__(self, index):

        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = self.vis_processor(image)

        return {
            "image": image,
            "image_id": ann["image_id"],
            "instance_id": ann["instance_id"],
        }


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import os
import random

from collections import OrderedDict

from lavis.datasets.datasets.multimodal_classification_datasets import (
    MultimodalClassificationDataset,
)
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file_L": ann["images"][0],
                "file_R": ann["images"][1],
                "sentence": ann["sentence"],
                "label": ann["label"],
                "image": [sample["image0"], sample["image1"]],
            }
        )

class NLVRDataset(MultimodalClassificationDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.class_labels = self._build_class_labels()

    def _build_class_labels(self):
        return {"False": 0, "True": 1}

    @staticmethod
    def _flip(samples):
        sentence = samples["text_input"]
        image0, image1 = samples["image0"], samples["image1"]

        if "left" not in sentence and "right" not in sentence:
            if random.random() < 0.5:
                image0, image1 = image1, image0
        else:
            if random.random() < 0.5:
                sentence = sentence.replace("left", "[TEMP_TOKEN]")
                sentence = sentence.replace("right", "left")
                sentence = sentence.replace("[TEMP_TOKEN]", "right")

                image0, image1 = image1, image0

        samples["text_input"] = sentence
        samples["image0"] = image0
        samples["image1"] = image1

        return samples

    def __getitem__(self, index):
        ann = self.annotation[index]

        image0_path = os.path.join(self.vis_root, ann["images"][0])
        image0 = Image.open(image0_path).convert("RGB")
        image0 = self.vis_processor(image0)

        image1_path = os.path.join(self.vis_root, ann["images"][1])
        image1 = Image.open(image1_path).convert("RGB")
        image1 = self.vis_processor(image1)

        sentence = self.text_processor(ann["sentence"])
        label = self.class_labels[ann["label"]]

        return self._flip(
            {
                "image0": image0,
                "image1": image1,
                "text_input": sentence,
                "label": label,
                # "image_id": ann["image_id"],
                "instance_id": ann["instance_id"],
            }
        )

class NLVREvalDataset(NLVRDataset):
    @staticmethod
    def _flip(samples):
        return samples


"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import os
import random

from collections import OrderedDict

from lavis.datasets.datasets.multimodal_classification_datasets import (
    MultimodalClassificationDataset,
)
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file_L": ann["images"][0],
                "file_R": ann["images"][1],
                "sentence": ann["sentence"],
                "label": ann["label"],
                "image": [sample["image0"], sample["image1"]],
            }
        )

class NLVRDataset(MultimodalClassificationDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.class_labels = self._build_class_labels()

    def _build_class_labels(self):
        return {"False": 0, "True": 1}

    @staticmethod
    def _flip(samples):
        sentence = samples["text_input"]
        image0, image1 = samples["image0"], samples["image1"]

        if "left" not in sentence and "right" not in sentence:
            if random.random() < 0.5:
                image0, image1 = image1, image0
        else:


NLVRDataset

    MultimodalClassificationDataset,
)
from PIL import Image

class __DisplMixin:
    def displ_item(self, index):
        sample, ann = self.__getitem__(index), self.annotation[index]

        return OrderedDict(
            {
                "file_L": ann["images"][0],
                "file_R": ann["images"][1],
                "sentence": ann["sentence"],
                "label": ann["label"],
                "image": [sample["image0"], sample["image1"]],
            }
        )

class NLVRDataset(MultimodalClassificationDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.class_labels = self._build_class_labels()

    def _build_class_labels(self):
        return {"False": 0, "True": 1}

    @staticmethod
    def _flip(samples):
        sentence = samples["text_input"]
        image0, image1 = samples["image0"], samples["image1"]

        if "left" not in sentence and "right" not in sentence:
            if random.random() < 0.5:
                image0, image1 = image1, image0
        else:
            if random.random() < 0.5:
                sentence = sentence.replace("left", "[TEMP_TOKEN]")
                sentence = sentence.replace("right", "left")
                sentence = sentence.replace("[TEMP_TOKEN]", "right")

                image0, image1 = image1, image0

        samples["text_input"] = sentence
        samples["image0"] = image0
        samples["image1"] = image1

        return samples

    def __getitem__(self, index):
        ann = self.annotation[index]

        image0_path = os.path.join(self.vis_root, ann["images"][0])
        image0 = Image.open(image0_path).convert("RGB")
        image0 = self.vis_processor(image0)

        image1_path = os.path.join(self.vis_root, ann["images"][1])
        image1 = Image.open(image1_path).convert("RGB")
        image1 = self.vis_processor(image1)

        sentence = self.text_processor(ann["sentence"])
        label = self.class_labels[ann["label"]]

        return self._flip(
            {
                "image0": image0,
                "image1": image1,
                "text_input": sentence,
                "label": label,
                # "image_id": ann["image_id"],
                "instance_id": ann["instance_id"],
            }
        )

class NLVREvalDataset(NLVRDataset):
    @staticmethod
    def _flip(samples):
        return samples