'''
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 '<PERSON>' in caption: # to handle with GCC12M
caption = caption.replace('<PERSON>', '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']")