Add GeoText inference endpoint support

4d3c316 18 days ago

11.8 kB

import os
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image

from dataset.re_bbox_dataset import re_dataset_bbox, re_eval_dataset

from dataset.randaugment import RandomAugment
#transform 去除了对于图像的水平旋转和镜像

'''
RandomResizedCrop:此操作首先随机调整图像的大小和纵横比,然后对其进行裁剪。结果图像的大小会是config['image_res']。scale=(0.2, 1.0)表示图像的原始面积与裁剪面积之间的比例范围,这里是0.2到1.0之间。interpolation=Image.BICUBIC表示在缩放图像时使用双三次插值。
RandomHorizontalFlip:以50%的概率对图像进行水平翻转。
RandomAugment:这是一种数据增强方法,但并不是torchvision库中的标准方法,可能是某个特定库或自定义实现的。从其参数可以看出,它将在给定的数据增强方法列表augs中随机选择2种方法,并应用最多7次的随机增强。所列出的增强方法包括:
Identity:保持不变
AutoContrast:自动调整图像的对比度
Equalize:使图像的直方图均衡化
Brightness:调整图像的亮度
Sharpness:调整图像的锐度
ShearX、ShearY:对图像进行X轴或Y轴的扭曲
TranslateX、TranslateY:对图像在X轴或Y轴上进行随机平移
Rotate 随机旋转图像
ToTensor 将PIL Image或numpy ndarray的图像数据转换为torch.Tensor 并且调整其范围从[0, 255]到[0.0, 1.0]。
normalize 对图像数据进行归一化。具体的均值和标准差应该在normalize变量中定义 但您没有提供这部分代码。归一化有助于神经网络的训练稳定和加速。
'''
def create_dataset(dataset, config, evaluate=False):
    normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))

    pretrain_transform = transforms.Compose([
        transforms.RandomResizedCrop(config['image_res'], scale=(0.2, 1.0),
                                     interpolation=Image.BICUBIC),
        # transforms.RandomHorizontalFlip(),
        RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
        transforms.ToTensor(),
        normalize,
    ])

    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(config['image_res'], scale=(0.5, 1.0),
                                     interpolation=Image.BICUBIC),
        # transforms.RandomHorizontalFlip(),
        RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
        transforms.ToTensor(),
        normalize,
    ])

    train_transform_wohflip = transforms.Compose([
        transforms.RandomResizedCrop(config['image_res'], scale=(0.5, 1.0),
                                     interpolation=Image.BICUBIC),
        # transforms.RandomHorizontalFlip(),
        RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness',]),
        transforms.ToTensor(),
        normalize,
    ])

    box_transform = transforms.Compose([
        RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
        transforms.ToTensor(),
        normalize,
    ])

    test_transform = transforms.Compose([
        transforms.Resize((config['image_res'], config['image_res']), interpolation=Image.BICUBIC),
        transforms.ToTensor(),
        normalize,
    ])

    if dataset == 'pretrain':
        general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
                                               world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
                                               transform=pretrain_transform)

        region_dataset = RegionTextJsonDataset(config, config['train_file_regions'], rank=int(os.environ.get('RANK') or 0),
                                                world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
                                                transform=pretrain_transform, box_transform=box_transform)

        return general_dataset, region_dataset

    elif dataset == 're':
        test_dataset = re_eval_dataset(config['test_file'], test_transform, config['image_root'])
        if evaluate:
            return None, test_dataset

        train_dataset = re_train_dataset(config['train_file'], train_transform, config['image_root'])
        
        # val_dataset = re_eval_dataset(config['val_file'], test_transform, config['image_root'])
        return train_dataset, test_dataset
    #修改

    elif dataset == 're_bbox':
        test_dataset = re_eval_dataset(config['test_file'], test_transform, config['image_root'])
        if evaluate:
            return None, test_dataset

        train_transform = transforms.Compose([
            RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
            transforms.ToTensor(),
            normalize,
        ])
        
        train_dataset = re_dataset_bbox(config['train_file'], train_transform, config['image_root'], mode='train', config=config)
        return train_dataset, test_dataset

    elif dataset == 'vqa':
        vqa_test_dataset = vqa_dataset(config['test_file'], test_transform, config['vqa_root'], config['vg_root'],
                                       split='test', answer_list=config['answer_list'],
                                       text_encoder=config['text_encoder'], use_roberta=config['use_roberta'])
        if evaluate:
            return None, vqa_test_dataset

        train_dataset = vqa_dataset(config['train_file'], train_transform_wohflip, config['vqa_root'], config['vg_root'],
                                    split='train', text_encoder=config['text_encoder'], use_roberta=config['use_roberta'])
        return train_dataset, vqa_test_dataset

    elif dataset == 'nlvr_pretrain':
        general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
                                               world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
                                               transform=pretrain_transform)

        return general_dataset

    elif dataset == 'nlvr':
        test_dataset = nlvr_dataset(config['test_file'], test_transform, config['image_root'])
        if evaluate:
            return None, None, test_dataset

        train_dataset = nlvr_dataset(config['train_file'], train_transform, config['image_root'])
        val_dataset = nlvr_dataset(config['val_file'], test_transform, config['image_root'])
        return train_dataset, val_dataset, test_dataset

    elif dataset == 'grounding':
        test_dataset = grounding_dataset(config['test_file'], test_transform, config['image_root'], mode='test')
        if evaluate:
            return None, test_dataset

        train_transform = transforms.Compose([
            transforms.Resize((config['image_res'], config['image_res']), interpolation=Image.BICUBIC),
            transforms.RandomHorizontalFlip(),
            RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
            transforms.ToTensor(),
            normalize,
        ])
        train_dataset = grounding_dataset(config['train_file'], train_transform, config['image_root'], mode='train')
        return train_dataset, test_dataset

    elif dataset == 'grounding_bbox_pretrain':
        region_dataset = RegionTextJsonDataset(config, config['train_file_regions'], rank=int(os.environ.get('RANK') or 0),
                                                world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
                                                transform=pretrain_transform, box_transform=box_transform)

        return region_dataset

    elif dataset == 'grounding_bbox':
        test_dataset = grounding_dataset_bbox(config['test_file'], test_transform, config['image_root'], mode='test', config=config)
        if evaluate:
            return None, test_dataset

        train_transform = transforms.Compose([
            RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
            transforms.ToTensor(),
            normalize,
        ])
        train_dataset = grounding_dataset_bbox(config['train_file'], train_transform, config['image_root'], mode='train', config=config)
        return train_dataset, test_dataset

    elif dataset == 'captioning_pretrain':
        general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
                                               world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
                                               transform=pretrain_transform, add_eos=True)
        return general_dataset

    elif dataset == 'caption_coco':
        train_dataset = coco_karpathy_train(train_transform, config['image_root'], config['train_file'], prompt=config['prompt'], max_words=config['max_tokens'])
        val_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['val_file'], 'val')
        test_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['test_file'], 'test')

        return train_dataset, val_dataset, test_dataset

    elif dataset == 'caption_coco_scst':
        train_dataset = coco_karpathy_train_scst(train_transform, config['image_root'], config['train_file'],
                                            prompt=config['prompt'], max_words=config['max_tokens'])
        val_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['val_file'], 'val')
        test_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['test_file'], 'test')

        return train_dataset, val_dataset, test_dataset

    else:
        raise NotImplementedError(f"dataset == {dataset}")


def vqa_collate_fn(batch):
    image_list, question_list, answer_list, weight_list, n = [], [], [], [], []
    for image, question, answer, weights in batch:
        image_list.append(image)
        question_list.append(question)
        weight_list += weights       
        answer_list += answer
        n.append(len(answer))
    return torch.stack(image_list, dim=0), question_list, answer_list, torch.Tensor(weight_list), n


def create_sampler(datasets, shuffles, num_tasks, global_rank):
    samplers = []
    for dataset, shuffle in zip(datasets, shuffles):
        sampler = torch.utils.data.DistributedSampler(dataset, num_replicas=num_tasks, rank=global_rank, shuffle=shuffle)
        samplers.append(sampler)
    return samplers


def create_loader(datasets, samplers, batch_size, num_workers, is_trains, collate_fns):
    loaders = []
    for dataset, sampler, bs, n_worker, is_train, collate_fn in zip(datasets, samplers, batch_size, num_workers,
                                                                    is_trains, collate_fns):
        if is_train:
            shuffle = (sampler is None)
            drop_last = True
        else:
            shuffle = False
            drop_last = False
        loader = DataLoader(
            dataset,
            batch_size=bs,
            num_workers=n_worker,
            pin_memory=True,
            sampler=sampler,
            shuffle=shuffle,
            collate_fn=collate_fn,
            drop_last=drop_last,
        )
        loaders.append(loader)

    if len(loaders) <= 1:
        print(f"### be careful: func create_loader returns a list length of {len(loaders)}")

    return loaders

	import os
	import torch
	from torch.utils.data import DataLoader
	from torchvision import transforms
	from PIL import Image

	from dataset.re_bbox_dataset import re_dataset_bbox, re_eval_dataset

	from dataset.randaugment import RandomAugment
	#transform 去除了对于图像的水平旋转和镜像

	'''
	RandomResizedCrop:此操作首先随机调整图像的大小和纵横比,然后对其进行裁剪。结果图像的大小会是config['image_res']。scale=(0.2, 1.0)表示图像的原始面积与裁剪面积之间的比例范围,这里是0.2到1.0之间。interpolation=Image.BICUBIC表示在缩放图像时使用双三次插值。

	RandomHorizontalFlip:以50%的概率对图像进行水平翻转。

	RandomAugment:这是一种数据增强方法,但并不是torchvision库中的标准方法,可能是某个特定库或自定义实现的。从其参数可以看出,它将在给定的数据增强方法列表augs中随机选择2种方法,并应用最多7次的随机增强。所列出的增强方法包括:

	Identity:保持不变
	AutoContrast:自动调整图像的对比度
	Equalize:使图像的直方图均衡化
	Brightness:调整图像的亮度
	Sharpness:调整图像的锐度
	ShearX、ShearY:对图像进行X轴或Y轴的扭曲
	TranslateX、TranslateY:对图像在X轴或Y轴上进行随机平移
	Rotate 随机旋转图像
	ToTensor 将PIL Image或numpy ndarray的图像数据转换为torch.Tensor 并且调整其范围从[0, 255]到[0.0, 1.0]。

	normalize 对图像数据进行归一化。具体的均值和标准差应该在normalize变量中定义但您没有提供这部分代码。归一化有助于神经网络的训练稳定和加速。
	'''
	def create_dataset(dataset, config, evaluate=False):
	normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))

	pretrain_transform = transforms.Compose([
	transforms.RandomResizedCrop(config['image_res'], scale=(0.2, 1.0),
	interpolation=Image.BICUBIC),
	# transforms.RandomHorizontalFlip(),
	RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
	transforms.ToTensor(),
	normalize,
	])

	train_transform = transforms.Compose([
	transforms.RandomResizedCrop(config['image_res'], scale=(0.5, 1.0),
	interpolation=Image.BICUBIC),
	# transforms.RandomHorizontalFlip(),
	RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
	transforms.ToTensor(),
	normalize,
	])

	train_transform_wohflip = transforms.Compose([
	transforms.RandomResizedCrop(config['image_res'], scale=(0.5, 1.0),
	interpolation=Image.BICUBIC),
	# transforms.RandomHorizontalFlip(),
	RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness',]),
	transforms.ToTensor(),
	normalize,
	])

	box_transform = transforms.Compose([
	RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
	transforms.ToTensor(),
	normalize,
	])

	test_transform = transforms.Compose([
	transforms.Resize((config['image_res'], config['image_res']), interpolation=Image.BICUBIC),
	transforms.ToTensor(),
	normalize,
	])

	if dataset == 'pretrain':
	general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
	world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
	transform=pretrain_transform)

	region_dataset = RegionTextJsonDataset(config, config['train_file_regions'], rank=int(os.environ.get('RANK') or 0),
	world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
	transform=pretrain_transform, box_transform=box_transform)

	return general_dataset, region_dataset

	elif dataset == 're':
	test_dataset = re_eval_dataset(config['test_file'], test_transform, config['image_root'])
	if evaluate:
	return None, test_dataset

	train_dataset = re_train_dataset(config['train_file'], train_transform, config['image_root'])

	# val_dataset = re_eval_dataset(config['val_file'], test_transform, config['image_root'])
	return train_dataset, test_dataset
	#修改

	elif dataset == 're_bbox':
	test_dataset = re_eval_dataset(config['test_file'], test_transform, config['image_root'])
	if evaluate:
	return None, test_dataset

	train_transform = transforms.Compose([
	RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
	transforms.ToTensor(),
	normalize,
	])

	train_dataset = re_dataset_bbox(config['train_file'], train_transform, config['image_root'], mode='train', config=config)
	return train_dataset, test_dataset

	elif dataset == 'vqa':
	vqa_test_dataset = vqa_dataset(config['test_file'], test_transform, config['vqa_root'], config['vg_root'],
	split='test', answer_list=config['answer_list'],
	text_encoder=config['text_encoder'], use_roberta=config['use_roberta'])
	if evaluate:
	return None, vqa_test_dataset

	train_dataset = vqa_dataset(config['train_file'], train_transform_wohflip, config['vqa_root'], config['vg_root'],
	split='train', text_encoder=config['text_encoder'], use_roberta=config['use_roberta'])
	return train_dataset, vqa_test_dataset

	elif dataset == 'nlvr_pretrain':
	general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
	world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
	transform=pretrain_transform)

	return general_dataset

	elif dataset == 'nlvr':
	test_dataset = nlvr_dataset(config['test_file'], test_transform, config['image_root'])
	if evaluate:
	return None, None, test_dataset

	train_dataset = nlvr_dataset(config['train_file'], train_transform, config['image_root'])
	val_dataset = nlvr_dataset(config['val_file'], test_transform, config['image_root'])
	return train_dataset, val_dataset, test_dataset

	elif dataset == 'grounding':
	test_dataset = grounding_dataset(config['test_file'], test_transform, config['image_root'], mode='test')
	if evaluate:
	return None, test_dataset

	train_transform = transforms.Compose([
	transforms.Resize((config['image_res'], config['image_res']), interpolation=Image.BICUBIC),
	transforms.RandomHorizontalFlip(),
	RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
	transforms.ToTensor(),
	normalize,
	])
	train_dataset = grounding_dataset(config['train_file'], train_transform, config['image_root'], mode='train')
	return train_dataset, test_dataset

	elif dataset == 'grounding_bbox_pretrain':
	region_dataset = RegionTextJsonDataset(config, config['train_file_regions'], rank=int(os.environ.get('RANK') or 0),
	world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
	transform=pretrain_transform, box_transform=box_transform)

	return region_dataset

	elif dataset == 'grounding_bbox':
	test_dataset = grounding_dataset_bbox(config['test_file'], test_transform, config['image_root'], mode='test', config=config)
	if evaluate:
	return None, test_dataset

	train_transform = transforms.Compose([
	RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness']),
	transforms.ToTensor(),
	normalize,
	])
	train_dataset = grounding_dataset_bbox(config['train_file'], train_transform, config['image_root'], mode='train', config=config)
	return train_dataset, test_dataset

	elif dataset == 'captioning_pretrain':
	general_dataset = ImageTextJsonDataset(config, config['train_file'], rank=int(os.environ.get('RANK') or 0),
	world_size=int(os.environ.get('WORLD_SIZE') or 1), shuffle=True, repeat=True,
	transform=pretrain_transform, add_eos=True)
	return general_dataset

	elif dataset == 'caption_coco':
	train_dataset = coco_karpathy_train(train_transform, config['image_root'], config['train_file'], prompt=config['prompt'], max_words=config['max_tokens'])
	val_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['val_file'], 'val')
	test_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['test_file'], 'test')

	return train_dataset, val_dataset, test_dataset

	elif dataset == 'caption_coco_scst':
	train_dataset = coco_karpathy_train_scst(train_transform, config['image_root'], config['train_file'],
	prompt=config['prompt'], max_words=config['max_tokens'])
	val_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['val_file'], 'val')
	test_dataset = coco_karpathy_caption_eval(test_transform, config['image_root'], config['test_file'], 'test')

	return train_dataset, val_dataset, test_dataset

	else:
	raise NotImplementedError(f"dataset == {dataset}")


	def vqa_collate_fn(batch):
	image_list, question_list, answer_list, weight_list, n = [], [], [], [], []
	for image, question, answer, weights in batch:
	image_list.append(image)
	question_list.append(question)
	weight_list += weights
	answer_list += answer
	n.append(len(answer))
	return torch.stack(image_list, dim=0), question_list, answer_list, torch.Tensor(weight_list), n


	def create_sampler(datasets, shuffles, num_tasks, global_rank):
	samplers = []
	for dataset, shuffle in zip(datasets, shuffles):
	sampler = torch.utils.data.DistributedSampler(dataset, num_replicas=num_tasks, rank=global_rank, shuffle=shuffle)
	samplers.append(sampler)
	return samplers


	def create_loader(datasets, samplers, batch_size, num_workers, is_trains, collate_fns):
	loaders = []
	for dataset, sampler, bs, n_worker, is_train, collate_fn in zip(datasets, samplers, batch_size, num_workers,
	is_trains, collate_fns):
	if is_train:
	shuffle = (sampler is None)
	drop_last = True
	else:
	shuffle = False
	drop_last = False
	loader = DataLoader(
	dataset,
	batch_size=bs,
	num_workers=n_worker,
	pin_memory=True,
	sampler=sampler,
	shuffle=shuffle,
	collate_fn=collate_fn,
	drop_last=drop_last,
	)
	loaders.append(loader)

	if len(loaders) <= 1:
	print(f"### be careful: func create_loader returns a list length of {len(loaders)}")

	return loaders