feat: xtd-10 dataset (#134)

* add xtd10 dataset

* fixes

* fix dataset dir

* add de, fr and jp

* fix

* add log

* fix

* update readme

Author: mertalev

.gitignore

@@ -53,6 +53,5 @@ features
 root
 cifar-100-*
 probe_benchmark/data
-datasets
 downloads
 .env

README.md

@@ -214,6 +214,10 @@ For Flickr-8k (zero-shot retrieval)
 
 - `clip_benchmark eval --model xlm-roberta-base-ViT-B-32 --pretrained laion5b_s13b_b90k --dataset=flickr8k --output=result.json --batch_size=64 --language=<LANG>`, where `<LANG>` can be among `en` (english), `zh` (chinese).
 
+For XTD-10 (zero-shot retrieval)
+
+- `clip_benchmark eval --model xlm-roberta-base-ViT-B-32 --pretrained laion5b_s13b_b90k --dataset=xtd10 --output=result.json --batch_size=64 --language=<LANG>`, where `<LANG>` can be among `es` (spanish), `it` (italian), `jp` (japanese), `ko` (korean), `pl` (polish), `ru` (russian), `tr` (Turkish), `zh` (chinese), `en` (english), `fr` (french), `de` (german).
+
 For [Crossmodal-3600](https://google.github.io/crossmodal-3600/) (zero-shot retrieval)
 
 - `clip_benchmark eval --model xlm-roberta-base-ViT-B-32 --pretrained laion5b_s13b_b90k --dataset=crossmodal3600 --output=result.json --batch_size=64 --language=<LANG>`, see supported languages [here](https://github.com/LAION-AI/CLIP_benchmark/blob/main/clip_benchmark/datasets/crossmodal3600.py#L9).

clip_benchmark/datasets/builder.py

@@ -283,6 +283,16 @@ def download_imagenet(r):
             crossmodal3600.create_annotation_file(root, language)
 
         ds = crossmodal3600.Crossmodal3600(root=root, ann_file=annotation_file, transform=transform, **kwargs)
+    elif dataset_name == 'xtd10':
+        from clip_benchmark.datasets import xtd10
+        if language not in xtd10.SUPPORTED_LANGUAGES:
+            raise ValueError("Unsupported language for xtd10:", language)
+
+        annotation_file = os.path.join(root, xtd10.OUTPUT_FILENAME_TEMPLATE.format(language))
+        if not os.path.exists(annotation_file):
+            xtd10.create_annotation_file(root, language)
+
+        ds = xtd10.XTD10(root=root, ann_file=annotation_file, transform=transform, **kwargs)
     elif dataset_name == 'xtd200':
         from clip_benchmark.datasets import xtd200
         if language not in xtd200.SUPPORTED_LANGUAGES:
@@ -523,15 +533,15 @@ def __len__(self):
         return 1
 
 def get_dataset_default_task(dataset):
-    if dataset in ("flickr30k", "flickr8k", "mscoco_captions", "multilingual_mscoco_captions", "flickr30k-200", "crossmodal3600", "xtd200"):
+    if dataset in ("flickr30k", "flickr8k", "mscoco_captions", "multilingual_mscoco_captions", "flickr30k-200", "crossmodal3600", "xtd10", "xtd200"):
         return "zeroshot_retrieval"
     elif dataset.startswith("sugar_crepe") or dataset == "winoground":
         return "image_caption_selection"
     else:
         return "zeroshot_classification"
 
 def get_dataset_collate_fn(dataset_name):
-    if dataset_name in ("mscoco_captions", "multilingual_mscoco_captions", "flickr30k", "flickr8k", "flickr30k-200", "crossmodal3600", "xtd200", "winoground") or dataset_name.startswith("sugar_crepe"):
+    if dataset_name in ("mscoco_captions", "multilingual_mscoco_captions", "flickr30k", "flickr8k", "flickr30k-200", "crossmodal3600", "xtd10", "xtd200", "winoground") or dataset_name.startswith("sugar_crepe"):
         return image_captions_collate_fn
     else:
         return default_collate

clip_benchmark/datasets/xtd10.py

@@ -0,0 +1,137 @@
+import codecs
+import json
+import os
+from subprocess import call
+
+import requests
+from PIL import Image
+from torchvision.datasets import VisionDataset
+
+
+GITHUB_DATA_PATH = "https://raw.githubusercontent.com/adobe-research/Cross-lingual-Test-Dataset-XTD10/main/XTD10"
+GITHUB_MIC_DATA_PATH = "https://raw.githubusercontent.com/adobe-research/Cross-lingual-Test-Dataset-XTD10/main/MIC"
+GITHUB_STAIR_DATA_PATH = "https://raw.githubusercontent.com/adobe-research/Cross-lingual-Test-Dataset-XTD10/main/STAIR"
+SUPPORTED_LANGUAGES = [
+    "de",
+    "en",
+    "es",
+    "fr",
+    "it",
+    "jp",
+    "ko",
+    "pl",
+    "ru",
+    "tr",
+    "zh",
+]
+
+IMAGE_INDEX_FILENAME = "test_image_names.txt"
+
+CAPTIONS_FILENAME_TEMPLATE = "test_1kcaptions_{}.txt"
+OUTPUT_FILENAME_TEMPLATE = "xtd10-{}.json"
+
+IMAGES_DOWNLOAD_URL = "https://nllb-data.com/test/xtd10/images.tar.gz"
+
+
+class XTD10(VisionDataset):
+    def __init__(self, root, ann_file, transform=None, target_transform=None):
+        super().__init__(root, transform=transform, target_transform=target_transform)
+        self.ann_file = os.path.expanduser(ann_file)
+        with codecs.open(ann_file, "r", encoding="utf-8") as fp:
+            data = json.load(fp)
+        self.data = [
+            (img_path, txt)
+            for img_path, txt in zip(data["image_paths"], data["annotations"])
+        ]
+
+    def __getitem__(self, index):
+        img, captions = self.data[index]
+
+        # Image
+        img = Image.open(img).convert("RGB")
+        if self.transform is not None:
+            img = self.transform(img)
+
+        # Captions
+        target = [
+            captions,
+        ]
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+
+def _get_lines(url):
+    response = requests.get(url, timeout=30)
+    return response.text.splitlines()
+
+
+def _download_images(out_path):
+    os.makedirs(out_path, exist_ok=True)
+    print("Downloading images")
+    call(f"wget {IMAGES_DOWNLOAD_URL} -O images.tar.gz", shell=True)
+    call(f"tar -xzf images.tar.gz -C {out_path}", shell=True)
+    call("rm images.tar.gz", shell=True)
+
+
+def create_annotation_file(root, lang_code):
+    if lang_code not in SUPPORTED_LANGUAGES:
+        raise ValueError(
+            f"Language code {lang_code} not supported. Supported languages are {SUPPORTED_LANGUAGES}"
+        )
+    data_dir = os.path.join(root, "xtd10")
+    if not os.path.exists(data_dir):
+        _download_images(data_dir)
+    images_dir = os.path.join(data_dir, "images")
+    print("Downloading xtd10 index file")
+    download_path = os.path.join(GITHUB_DATA_PATH, IMAGE_INDEX_FILENAME)
+    target_images = _get_lines(download_path)
+
+    print("Downloading xtd10 captions:", lang_code)
+    captions_path = GITHUB_DATA_PATH
+    match lang_code:
+        case "de" | "fr":
+            captions_path = GITHUB_MIC_DATA_PATH
+        case "jp":
+            captions_path = GITHUB_STAIR_DATA_PATH
+        case _:
+            captions_path = GITHUB_DATA_PATH
+    download_path = os.path.join(
+        captions_path, CAPTIONS_FILENAME_TEMPLATE.format(lang_code)
+    )
+    target_captions = _get_lines(download_path)
+
+    number_of_missing_images = 0
+    valid_images, valid_annotations, valid_indicies = [], [], []
+    for i, (img, txt) in enumerate(zip(target_images, target_captions)):
+        image_path = os.path.join(images_dir, img)
+        if not os.path.exists(image_path):
+            print("Missing image file", img)
+            number_of_missing_images += 1
+            continue
+
+        valid_images.append(image_path)
+        valid_annotations.append(txt)
+        valid_indicies.append(i)
+
+    if number_of_missing_images > 0:
+        print(f"*** WARNING *** missing {number_of_missing_images} files.")
+
+    with codecs.open(
+        os.path.join(root, OUTPUT_FILENAME_TEMPLATE.format(lang_code)),
+        "w",
+        encoding="utf-8",
+    ) as fp:
+        json.dump(
+            {
+                "image_paths": valid_images,
+                "annotations": valid_annotations,
+                "indicies": valid_indicies,
+            },
+            fp,
+            ensure_ascii=False,
+        )