''' LinCIR Copyright (c) 2023-present NAVER Corp. CC BY-NC-4.0 (https://creativecommons.org/licenses/by-nc/4.0/) ''' import os import functools import glob import random import json from pathlib import Path from typing import List, Optional, Union, Dict, Literal import PIL import PIL.Image import torch from torch.utils.data import Dataset import webdataset as wds import spacy import numpy as np import sng_parser import datasets def extract_keywords(spacy_nlp, caption): candidates = [] nlp_caption = caption doc = spacy_nlp(nlp_caption) tmp = '' for word in doc: if word.pos_ == 'ADJ': if tmp == '': tmp += word.text else: tmp += ' ' + word.text elif word.pos_ == 'NOUN' or word.pos_ == 'PROPN': if tmp == '': tmp += word.text else: tmp += ' ' + word.text else: if tmp != '': candidates.append(tmp) tmp = '' if tmp != '': candidates.append(tmp) candidates = list(set(candidates)) return candidates def extract_keywords_spacy(spacy_nlp, caption): sequences = [] current_sequence = [] doc = spacy_nlp(caption) for token in doc: # Check if the token is a noun, proper noun, or adjective if token.pos_ in ['NOUN', 'PROPN', 'ADJ', 'DET']: current_sequence.append(token.text) else: # If we encounter a token that's not one of the desired POS and current_sequence is not empty if current_sequence: sequences.append(" ".join(current_sequence)) current_sequence = [] # Adding any remaining sequence after the loop if current_sequence: sequences.append(" ".join(current_sequence)) return sequences def extract_sng(caption): graph = sng_parser.parse(caption) entities = [x['head'] for i, x in enumerate(graph['entities'])] relations = [{'subject': entities[x['subject']], 'object': entities[x['object']], 'relation': x['relation']} for x in graph['relations']] return entities, relations def clean_caption(caption, tokenizer): if caption is None: caption = '' if '' in caption: # to handle with GCC12M caption = caption.replace('', 'person') caption = caption.lower().replace('$', '').strip() tokens = tokenizer.encode(caption, padding='longest', return_tensors='pt') if tokens.shape[1] > 77: caption = tokenizer.batch_decode(tokens[:,1:76])[0] return caption def preprocess_precomputed_base(sample, spacy_nlp, keywords_list, tokenizer): ''' 'image_feature.npy','json' ''' image_feature, image_feature_giga, meta = sample caption = clean_caption(meta['source_caption'], tokenizer) keywords = [''] try: keywords = extract_keywords_spacy(spacy_nlp, caption) except Exception as e: #print(e) pass # for keywords indicator = 1 replaced_caption = caption for keyword in keywords: if keyword != '' and keyword in caption: replaced_caption = replaced_caption.replace(keyword, '[$]') else: tmp_keywords = caption.split(' ') if len(tmp_keywords) > 0: selected_keywords = random.sample(tmp_keywords, k=min(int(len(tmp_keywords) * 1.0), 1)) for selected_keyword in selected_keywords: replaced_caption = replaced_caption.replace(selected_keyword, '[$]') else: replaced_caption = f'a photo of [$] that {caption}' indicator = 0 break token_dict = tokenizer(text=caption, return_tensors='pt', padding='max_length', truncation=True) tokens, attention_mask = token_dict['input_ids'][0], token_dict['attention_mask'][0] replaced_token_dict = tokenizer(text=replaced_caption, return_tensors='pt', padding='max_length', truncation=True) replaced_tokens, replaced_attention_mask = replaced_token_dict['input_ids'][0], replaced_token_dict['attention_mask'][0] replaced_tokens = torch.where(replaced_tokens == 49408, torch.ones_like(replaced_tokens) * 259, replaced_tokens) if 259 not in replaced_tokens: replaced_caption = 'a photo of [$]' replaced_token_dict = tokenizer(text=replaced_caption, return_tensors='pt', padding='max_length', truncation=True) replaced_tokens, replaced_attention_mask = replaced_token_dict['input_ids'][0], replaced_token_dict['attention_mask'][0] replaced_tokens = torch.where(replaced_tokens == 49408, torch.ones_like(replaced_tokens) * 259, replaced_tokens) indicator = 0 new_sample = [tokens, replaced_tokens, indicator] return tuple(new_sample) class CaptionDataset(Dataset): def __init__(self, captions, tokenizer, spacy_nlp): self.captions = captions self.tokenizer = tokenizer self.spacy_nlp = spacy_nlp def __len__(self): return len(self.captions) def __getitem__(self, idx): caption = self.captions[idx] caption = clean_caption(caption, self.tokenizer) keywords = [""] try: keywords = extract_keywords_spacy(self.spacy_nlp, caption) except Exception as e: #print(e) pass # for keywords indicator = 1 replaced_caption = caption if len(keywords) == 0: keywords = [""] for keyword in keywords: if keyword != '' and keyword in caption: replaced_caption = replaced_caption.replace(keyword, '[$]') else: tmp_keywords = caption.split(' ') if len(tmp_keywords) > 0: selected_keywords = random.sample(tmp_keywords, k=min(int(len(tmp_keywords) * 1.0), 1)) for selected_keyword in selected_keywords: replaced_caption = replaced_caption.replace(selected_keyword, '[$]') else: replaced_caption = f'a photo of [$] that {caption}' indicator = 0 break token_dict = self.tokenizer(text=caption, return_tensors='pt', padding='max_length', truncation=True) tokens, attention_mask = token_dict['input_ids'][0], token_dict['attention_mask'][0] replaced_token_dict = self.tokenizer(text=replaced_caption, return_tensors='pt', padding='max_length', truncation=True) replaced_tokens, replaced_attention_mask = replaced_token_dict['input_ids'][0], replaced_token_dict['attention_mask'][0] replaced_tokens = torch.where(replaced_tokens == 49408, torch.ones_like(replaced_tokens) * 259, replaced_tokens) if 259 not in replaced_tokens: replaced_caption = 'a photo of [$]' replaced_token_dict = self.tokenizer(text=replaced_caption, return_tensors='pt', padding='max_length', truncation=True) replaced_tokens, replaced_attention_mask = replaced_token_dict['input_ids'][0], replaced_token_dict['attention_mask'][0] replaced_tokens = torch.where(replaced_tokens == 49408, torch.ones_like(replaced_tokens) * 259, replaced_tokens) indicator = 0 return tokens, replaced_tokens, indicator def build_loader(args, tokenizer, accelerator): data_names = {'dataset1': 'dangne/gcc_caption_only', 'dataset2': 'FredZhang7/stable-diffusion-prompts-2.47M', 'dataset3': 'Geonmo/midjourney-prompts-only', } for k, v in data_names.items(): if not os.path.exists(os.path.join('./datasets', k)): if accelerator.is_main_process: print('Downloading captions is required') db = datasets.load_dataset(v, cache_dir=os.path.join('./datasets', k)) captions = [] for k, v in data_names.items(): db = datasets.load_dataset(v, cache_dir=os.path.join('./datasets', k)) captions += db['train']['text'] dataset = CaptionDataset(captions, tokenizer, spacy.load('en_core_web_sm')) data_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, num_workers=args.num_workers, drop_last=True, shuffle=True) return data_loader class FashionIQDataset(Dataset): """ Copy-paste from https://github.com/miccunifi/SEARLE/blob/main/src/datasets.py FashionIQ dataset class for PyTorch. The dataset can be used in 'relative' or 'classic' mode: - In 'classic' mode the dataset yield :a dict with keys ['image', 'image_name'] - In 'relative' mode the dataset yield dict with keys: - ['reference_image', 'reference_name', 'target_image', 'target_name', 'relative_captions'] when split in ['train', 'val'] - ['reference_image', 'reference_name', 'relative_captions'] when split == test """ def __init__(self, dataset_path: Union[Path, str], split: Literal['train', 'val', 'test'], dress_types: List[str], mode: Literal['relative', 'classic'], preprocess: callable, no_duplicates: Optional[bool] = False): """ :param dataset_path: path to the FashionIQ dataset :param split: dataset split, should be in ['train, 'val', 'test'] :param dress_types: list of fashionIQ categories, each category should be in ['dress', 'shirt', 'toptee'] :param mode: dataset mode, should be in ['relative', 'classic']: - In 'classic' mode the dataset yield a dict with keys ['image', 'image_name'] - In 'relative' mode the dataset yield dict with keys: - ['reference_image', 'reference_name', 'target_image', 'target_name', 'relative_captions'] when split in ['train', 'val'] - ['reference_image', 'reference_name', 'relative_captions'] when split == test :param preprocess: function which preprocesses the image :param no_duplicates: if True, the dataset will not yield duplicate images in relative mode, does not affect classic mode """ dataset_path = Path(dataset_path) self.dataset_path = dataset_path self.mode = mode self.dress_types = dress_types self.split = split self.no_duplicates = no_duplicates # Validate the inputs if mode not in ['relative', 'classic']: raise ValueError("mode should be in ['relative', 'classic']") if split not in ['test', 'train', 'val']: raise ValueError("split should be in ['test', 'train', 'val']") for dress_type in dress_types: if dress_type not in ['dress', 'shirt', 'toptee']: raise ValueError("dress_type should be in ['dress', 'shirt', 'toptee']") self.preprocess = preprocess # get triplets made by (reference_image, target_image, a pair of relative captions) self.triplets: List[dict] = [] for dress_type in dress_types: with open(dataset_path / 'captions' / f'cap.{dress_type}.{split}.json') as f: self.triplets.extend(json.load(f)) # Remove duplicats from if self.no_duplicates: seen = set() new_triplets = [] for triplet in self.triplets: if triplet['candidate'] not in seen: seen.add(triplet['candidate']) new_triplets.append(triplet) self.triplets = new_triplets # get the image names self.image_names: list = [] for dress_type in dress_types: with open(dataset_path / 'image_splits' / f'split.{dress_type}.{split}.json') as f: self.image_names.extend(json.load(f)) print(f"FashionIQ {split} - {dress_types} dataset in {mode} mode initialized") def __getitem__(self, index) -> dict: try: if self.mode == 'relative': relative_captions = self.triplets[index]['captions'] reference_name = self.triplets[index]['candidate'] if self.split in ['train', 'val']: reference_image_path = self.dataset_path / 'images' / f"{reference_name}.jpg" reference_image = self.preprocess(PIL.Image.open(reference_image_path), return_tensors='pt')['pixel_values'][0] target_name = self.triplets[index]['target'] target_image_path = self.dataset_path / 'images' / f"{target_name}.jpg" target_image = self.preprocess(PIL.Image.open(target_image_path), return_tensors='pt')['pixel_values'][0] return { 'reference_image': reference_image, 'reference_name': reference_name, 'target_image': target_image, 'target_name': target_name, 'relative_captions': relative_captions } elif self.split == 'test': reference_image_path = self.dataset_path / 'images' / f"{reference_name}.jpg" reference_image = self.preprocess(PIL.Image.open(reference_image_path), return_tensors='pt')['pixel_values'][0] return { 'reference_image': reference_image, 'reference_name': reference_name, 'relative_captions': relative_captions } elif self.mode == 'classic': image_name = self.image_names[index] image_path = self.dataset_path / 'images' / f"{image_name}.jpg" image = self.preprocess(PIL.Image.open(image_path), return_tensors='pt')['pixel_values'][0] return { 'image': image, 'image_name': image_name } else: raise ValueError("mode should be in ['relative', 'classic']") except Exception as e: print(f"Exception: {e}") def __len__(self): if self.mode == 'relative': return len(self.triplets) elif self.mode == 'classic': return len(self.image_names) else: raise ValueError("mode should be in ['relative', 'classic']") class CIRRDataset(Dataset): """ Copy-paste from https://github.com/miccunifi/SEARLE/blob/main/src/datasets.py CIRR dataset class for PyTorch dataloader. The dataset can be used in 'relative' or 'classic' mode: - In 'classic' mode the dataset yield a dict with keys ['image', 'image_name'] - In 'relative' mode the dataset yield dict with keys: - ['reference_image', 'reference_name', 'target_image', 'target_name', 'relative_caption', 'group_members'] when split in ['train', 'val'] - ['reference_image', 'reference_name' 'relative_caption', 'group_members', 'pair_id'] when split == test """ def __init__(self, dataset_path: Union[Path, str], split: Literal['train', 'val', 'test'], mode: Literal['relative', 'classic'], preprocess: callable, no_duplicates: Optional[bool] = False): """ :param dataset_path: path to the CIRR dataset :param split: dataset split, should be in ['train', 'val', 'test'] :param mode: dataset mode, should be in ['relative', 'classic']: - In 'classic' mode the dataset yield a dict with keys ['image', 'image_name'] - In 'relative' mode the dataset yield dict with keys: - ['reference_image', 'reference_name', 'target_image', 'target_name', 'relative_caption', 'group_members'] when split in ['train', 'val'] - ['reference_image', 'reference_name' 'relative_caption', 'group_members', 'pair_id'] when split == test :param preprocess: function which preprocesses the image :param no_duplicates: if True, the dataset will not yield duplicate images in relative mode, does not affect classic mode """ dataset_path = Path(dataset_path) self.dataset_path = dataset_path self.preprocess = preprocess self.mode = mode self.split = split self.no_duplicates = no_duplicates if split == "test": split = "test1" self.split = "test1" # Validate inputs if split not in ['test1', 'train', 'val']: raise ValueError("split should be in ['test1', 'train', 'val']") if mode not in ['relative', 'classic']: raise ValueError("mode should be in ['relative', 'classic']") # get triplets made by (reference_image, target_image, relative caption) with open(dataset_path / 'cirr' / 'captions' / f'cap.rc2.{split}.json') as f: self.triplets = json.load(f) # Remove duplicates from triplets if self.no_duplicates: seen = set() new_triplets = [] for triplet in self.triplets: if triplet['reference'] not in seen: seen.add(triplet['reference']) new_triplets.append(triplet) self.triplets = new_triplets # get a mapping from image name to relative path with open(dataset_path / 'cirr' / 'image_splits' / f'split.rc2.{split}.json') as f: self.name_to_relpath = json.load(f) print(f"CIRR {split} dataset in {mode} mode initialized") def __getitem__(self, index) -> dict: try: if self.mode == 'relative': group_members = self.triplets[index]['img_set']['members'] reference_name = self.triplets[index]['reference'] relative_caption = self.triplets[index]['caption'] if self.split in ['train', 'val']: reference_image_path = self.dataset_path / self.name_to_relpath[reference_name] reference_image = self.preprocess(PIL.Image.open(reference_image_path), return_tensors='pt')['pixel_values'][0] target_hard_name = self.triplets[index]['target_hard'] target_image_path = self.dataset_path / self.name_to_relpath[target_hard_name] target_image = self.preprocess(PIL.Image.open(target_image_path), return_tensors='pt')['pixel_values'][0] return { 'reference_image': reference_image, 'reference_name': reference_name, 'target_image': target_image, 'target_name': target_hard_name, 'relative_caption': relative_caption, 'group_members': group_members } elif self.split == 'test1': pair_id = self.triplets[index]['pairid'] reference_image_path = self.dataset_path / self.name_to_relpath[reference_name] reference_image = self.preprocess(PIL.Image.open(reference_image_path), return_tensors='pt')['pixel_values'][0] return { 'reference_image': reference_image, 'reference_name': reference_name, 'relative_caption': relative_caption, 'group_members': group_members, 'pair_id': pair_id } elif self.mode == 'classic': image_name = list(self.name_to_relpath.keys())[index] image_path = self.dataset_path / self.name_to_relpath[image_name] im = PIL.Image.open(image_path) image = self.preprocess(im, return_tensors='pt')['pixel_values'][0] return { 'image': image, 'image_name': image_name } else: raise ValueError("mode should be in ['relative', 'classic']") except Exception as e: print(f"Exception: {e}") def __len__(self): if self.mode == 'relative': return len(self.triplets) elif self.mode == 'classic': return len(self.name_to_relpath) else: raise ValueError("mode should be in ['relative', 'classic']") class CIRCODataset(Dataset): """ Copy-paste from https://github.com/miccunifi/SEARLE/blob/main/src/datasets.py CIRCO dataset class for PyTorch. The dataset can be used in 'relative' or 'classic' mode: - In 'classic' mode the dataset yield a dict with keys ['image', 'image_name'] - In 'relative' mode the dataset yield dict with keys: - ['reference_image', 'reference_name', 'target_image', 'target_name', 'relative_captions', 'shared_concept', 'gt_img_ids', 'query_id'] when split == 'val' - ['reference_image', 'reference_name', 'relative_captions', 'shared_concept', 'query_id'] when split == test """ def __init__(self, dataset_path: Union[str, Path], split: Literal['val', 'test'], mode: Literal['relative', 'classic'], preprocess: callable): """ Args: dataset_path (Union[str, Path]): path to CIRCO dataset split (str): dataset split, should be in ['test', 'val'] mode (str): dataset mode, should be in ['relative', 'classic'] preprocess (callable): function which preprocesses the image """ # Set dataset paths and configurations dataset_path = Path(dataset_path) self.mode = mode self.split = split self.preprocess = preprocess self.data_path = dataset_path # Ensure input arguments are valid if mode not in ['relative', 'classic']: raise ValueError("mode should be in ['relative', 'classic']") if split not in ['test', 'val']: raise ValueError("split should be in ['test', 'val']") # Load COCO images information with open(dataset_path / 'COCO2017_unlabeled' / "annotations" / "image_info_unlabeled2017.json", "r") as f: imgs_info = json.load(f) self.img_paths = [dataset_path / 'COCO2017_unlabeled' / "unlabeled2017" / img_info["file_name"] for img_info in imgs_info["images"]] self.img_ids = [img_info["id"] for img_info in imgs_info["images"]] self.img_ids_indexes_map = {str(img_id): i for i, img_id in enumerate(self.img_ids)} # get CIRCO annotations with open(dataset_path / 'annotations' / f'{split}.json', "r") as f: self.annotations: List[dict] = json.load(f) # Get maximum number of ground truth images (for padding when loading the images) self.max_num_gts = 23 # Maximum number of ground truth images print(f"CIRCODataset {split} dataset in {mode} mode initialized") def get_target_img_ids(self, index) -> Dict[str, int]: """ Returns the id of the target image and ground truth images for a given query Args: index (int): id of the query Returns: Dict[str, int]: dictionary containing target image id and a list of ground truth image ids """ return { 'target_img_id': self.annotations[index]['target_img_id'], 'gt_img_ids': self.annotations[index]['gt_img_ids'] } def __getitem__(self, index) -> dict: """ Returns a specific item from the dataset based on the index. In 'classic' mode, the dataset yields a dictionary with the following keys: [img, img_id] In 'relative' mode, the dataset yields dictionaries with the following keys: - [reference_img, reference_img_id, target_img, target_img_id, relative_caption, shared_concept, gt_img_ids, query_id] if split == val - [reference_img, reference_img_id, relative_caption, shared_concept, query_id] if split == test """ if self.mode == 'relative': # Get the query id query_id = str(self.annotations[index]['id']) # Get relative caption and shared concept relative_caption = self.annotations[index]['relative_caption'] shared_concept = self.annotations[index]['shared_concept'] # Get the reference image reference_img_id = str(self.annotations[index]['reference_img_id']) reference_img_path = self.img_paths[self.img_ids_indexes_map[reference_img_id]] reference_img = self.preprocess(PIL.Image.open(reference_img_path), return_tensors='pt')['pixel_values'][0] if self.split == 'val': # Get the target image and ground truth images target_img_id = str(self.annotations[index]['target_img_id']) gt_img_ids = [str(x) for x in self.annotations[index]['gt_img_ids']] target_img_path = self.img_paths[self.img_ids_indexes_map[target_img_id]] target_img = self.preprocess(PIL.Image.open(target_img_path), return_tensors='pt')['pixel_values'][0] # Pad ground truth image IDs with zeros for collate_fn gt_img_ids += [''] * (self.max_num_gts - len(gt_img_ids)) return { 'reference_image': reference_img, 'reference_name': reference_img_id, 'target_image': target_img, 'target_name': target_img_id, 'relative_caption': relative_caption, 'shared_concept': shared_concept, 'gt_img_ids': gt_img_ids, 'query_id': query_id, } elif self.split == 'test': return { 'reference_image': reference_img, 'reference_name': reference_img_id, 'relative_caption': relative_caption, 'shared_concept': shared_concept, 'query_id': query_id, } elif self.mode == 'classic': # Get image ID and image path img_id = str(self.img_ids[index]) img_path = self.img_paths[index] # Preprocess image and return img = self.preprocess(PIL.Image.open(img_path), return_tensors='pt')['pixel_values'][0] return { 'image': img, 'image_name': img_id } def __len__(self): """ Returns the length of the dataset. """ if self.mode == 'relative': return len(self.annotations) elif self.mode == 'classic': return len(self.img_ids) else: raise ValueError("mode should be in ['relative', 'classic']")