| 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_DATA_PATH_DE_FR = "https://raw.githubusercontent.com/adobe-research/Cross-lingual-Test-Dataset-XTD10/main/MIC/" |
| GITHUB_DATA_PATH_JP = "https://raw.githubusercontent.com/adobe-research/Cross-lingual-Test-Dataset-XTD10/main/STAIR/" |
| SUPPORTED_LANGUAGES = ["es", "it", "ko", "pl", "ru", "tr", "zh", "en", "de", "fr", "jp"] |
|
|
| IMAGE_INDEX_FILENAME = "test_image_names.txt" |
|
|
| CAPTIONS_FILENAME_TEMPLATE = "test_1kcaptions_{}.txt" |
| OUTPUT_FILENAME_TEMPLATE = "multilingual_mscoco_captions-{}.json" |
|
|
| IMAGES_DOWNLOAD_URL = "https://nllb-data.com/test/xtd10/images.tar.gz" |
|
|
|
|
| class Multilingual_MSCOCO(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] |
|
|
| |
| img = Image.open(img).convert("RGB") |
| if self.transform is not None: |
| img = self.transform(img) |
|
|
| |
| 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, "multilingual_mscoco") |
| if not os.path.exists(data_dir): |
| _download_images(data_dir) |
| images_dir = os.path.join(data_dir, "images") |
| print("Downloading multilingual_ms_coco index file") |
| download_path = os.path.join(GITHUB_DATA_PATH, IMAGE_INDEX_FILENAME) |
| target_images = _get_lines(download_path) |
|
|
| print("Downloading multilingual_ms_coco captions:", lang_code) |
| captions_path = GITHUB_DATA_PATH |
| if lang_code in ["de", "fr"]: |
| captions_path = GITHUB_DATA_PATH_DE_FR |
| elif lang_code == "jp": |
| captions_path = GITHUB_DATA_PATH_JP |
| 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, |
| ) |
|
|
