pretrain_ref_coco_sparse_embeddings.py

# 使用Refcoco/Refcoco+/Refcocog
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
import torch.utils.data
import torchvision

import os
from tqdm import tqdm
import argparse
import warnings
import logging
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import util.misc as utils
import datasets.samplers as samplers
import torch.distributed as dist
from einops import rearrange
from datasets.coco_eval import CocoEvaluator
from pycocotools import mask as coco_mask
from collections import namedtuple
import matplotlib.pyplot as plt
from PIL import Image

import opts
import refersam
import loss

import random

seed = 42
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
warnings.filterwarnings('ignore')


class ModulatedDetection(torchvision.datasets.CocoDetection):
    def __init__(self, img_folder, ann_file, image_set, return_masks):
        super(ModulatedDetection, self).__init__(img_folder, ann_file)
        self.prepare = ConvertCocoPolysToMask(return_masks)
        self.image_set = image_set

    def __getitem__(self, idx):
        instance_check = False
        while not instance_check:
            img, target = super(ModulatedDetection, self).__getitem__(idx)
            image_id = self.ids[idx]
            coco_img = self.coco.loadImgs(image_id)[0]
            caption = coco_img["caption"]
            dataset_name = coco_img["dataset_name"] if "dataset_name" in coco_img else None
            target = {"image_id": image_id, "annotations": target, "caption": caption}
            img, target = self.prepare(img, target)
            target["dataset_name"] = dataset_name
            for extra_key in ["sentence_id", "original_img_id", "original_id", "task_id"]:
                if extra_key in coco_img:
                    target[extra_key] = coco_img[extra_key] # box xyxy -> cxcywh
            # FIXME: handle "valid", since some box may be removed due to random crop
            target["valid"] = torch.tensor([1]) if len(target["area"]) != 0 else torch.tensor([0])

            if torch.any(target['valid'] == 1):  # at leatst one instance
                instance_check = True
            else:
                idx = random.randint(0, self.__len__() - 1)

            if self.image_set == "train":
                final_scales = [288, 320, 352, 392, 416, 448, 480, 512]
                final_scales_2 = [400, 500, 600]

                # reshape image and padding
                transform = transforms.Compose([
                    transforms.ToTensor(),
                ])
                transform3_1 = transforms.ColorJitter(brightness=(0.5, 1.5))
                transform3_2 = transforms.ColorJitter(contrast=(0.5, 1.5))
                transform3_3 = transforms.ColorJitter(saturation=(0.5, 1.5))
                transform3_4 = transforms.ColorJitter(hue=(-0.1, 0.1))
                transform4 = transforms.RandomHorizontalFlip(p=1)

                resized_images = transform(img) # [3 640 640]
                resized_masks = target['masks']

                random_h = random.choice(final_scales)
                random_w = random.choice(final_scales)
                random_h_2 = random.choice(final_scales_2)
                random_w_2 = random.choice(final_scales_2)
                # transform5 = transforms.RandomResizedCrop((random_h, random_w), scale=(0.5, 1.5))
                transform5 = transforms.Resize((random_h, random_w), interpolation=Image.NEAREST)
                transform5_2 = transforms.Resize((random_h_2, random_w_2), interpolation=Image.NEAREST)

                transform6 = transforms.RandomResizedCrop(size=(384, 600), scale=(0.8, 1.0), ratio=(0.75, 1.333))

                # 颜色增强
                if random.random() > 0.3:
                    # 随机调整亮度
                    resized_images = transform3_1(resized_images)
                if random.random() > 0.3:
                    # 随机调整对比度
                    resized_images = transform3_2(resized_images)
                if random.random() > 0.3:
                    # 随机调整饱和度
                    resized_images = transform3_3(resized_images)
                if random.random() > 0.3:
                    # 随机调整色调
                    resized_images = transform3_4(resized_images)
                if random.random() > 0.3:
                    # 随机调整对比度
                    resized_images = transform3_2(resized_images)

                # 随机翻转
                if random.random() > 0.5:
                    resized_images = transform4(resized_images)
                    resized_masks = transform4(resized_masks)
                    target['caption'] = caption.replace('left', '@').replace('right', 'left').replace('@', 'right')
                
                # 随机大小变化
                if random.random() > 0.5:
                    resized_images = transform5(resized_images)
                    resized_masks = transform5(resized_masks)
                else:
                    resized_images = transform5_2(resized_images)
                    resized_masks = transform5_2(resized_masks)
                    resized_images = transform6(resized_images)
                    resized_masks = transform6(resized_masks)
                    resized_images = transform5(resized_images)
                    resized_masks = transform5(resized_masks)

                # t5 = transforms.Resize((max_height, max_width))
                # resized_images = t5(resized_images)
                # resized_masks = t5(resized_masks)

                resized_images = resized_images.permute(1, 2, 0)
                resized_masks = resized_masks.squeeze(0)
            else:
                # reshape image and padding
                final_scales = [296, 328, 360, 392, 416, 448, 480, 512]
                transform = transforms.Compose([
                    transforms.ToTensor(),
                ])
                random_h = random.choice(final_scales)
                random_w = random.choice(final_scales)
                transform5 = transforms.Resize((random_h, random_w), interpolation=Image.NEAREST)

                resized_images = transform(img)
                resized_masks = target['masks']
                # 随机大小变化
                if random.random() > 0.5:
                    resized_images = transform5(resized_images)
                    resized_masks = transform5(resized_masks)

                resized_images = resized_images.permute(1, 2, 0) # [640 640 3]
                resized_masks = resized_masks.squeeze(0) # [640 640]

            target['masks'] = resized_masks

        return resized_images, target


def convert_coco_poly_to_mask(segmentations, height, width):
    masks = []
    for polygons in segmentations:
        # 将多边形分割转换为COCO格式的RLE编码
        rles = coco_mask.frPyObjects(polygons, height, width)
        mask = coco_mask.decode(rles)
        if len(mask.shape) < 3:
            mask = mask[..., None]
        mask = torch.as_tensor(mask, dtype=torch.uint8)
        mask = mask.any(dim=2)
        masks.append(mask)
    if masks:
        masks = torch.stack(masks, dim=0)
    else:
        masks = torch.zeros((0, height, width), dtype=torch.uint8)
    return masks


class ConvertCocoPolysToMask(object):
    def __init__(self, return_masks=False):
        self.return_masks = return_masks

    def __call__(self, image, target):
        w, h = image.size

        image_id = target["image_id"]
        image_id = torch.tensor([image_id])

        anno = target["annotations"]
        caption = target["caption"] if "caption" in target else None

        anno = [obj for obj in anno if "iscrowd" not in obj or obj["iscrowd"] == 0]

        boxes = [obj["bbox"] for obj in anno]
        # guard against no boxes via resizing
        boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
        boxes[:, 2:] += boxes[:, :2] # xminyminwh -> xyxy
        boxes[:, 0::2].clamp_(min=0, max=w)
        boxes[:, 1::2].clamp_(min=0, max=h)

        classes = [obj["category_id"] for obj in anno]
        classes = torch.tensor(classes, dtype=torch.int64)

        if self.return_masks:
            segmentations = [obj["segmentation"] for obj in anno]
            masks = convert_coco_poly_to_mask(segmentations, h, w)

        # keep the valid boxes
        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
        boxes = boxes[keep]
        classes = classes[keep]
        if self.return_masks:
            masks = masks[keep]

        target = {}
        target["boxes"] = boxes
        target["labels"] = classes
        if caption is not None:
            target["caption"] = caption
        if self.return_masks:
            target["masks"] = masks
        target["image_id"] = image_id

        # for conversion to coco api
        area = torch.tensor([obj["area"] for obj in anno])
        iscrowd = torch.tensor([obj["iscrowd"] if "iscrowd" in obj else 0 for obj in anno])
        target["area"] = area[keep]
        target["iscrowd"] = iscrowd[keep]
        target["valid"] = torch.tensor([1])
        target["orig_size"] = torch.as_tensor([int(h), int(w)])
        target["size"] = torch.as_tensor([int(h), int(w)])
        return image, target


def build_refexp(dataset_file, image_set):
    root = Path("coco")
    assert root.exists(), f"provided COCO path {root} does not exist"
    mode = "instances"
    dataset = dataset_file
    PATHS = {
        "train": (root / "train2014", root / dataset / f"{mode}_{dataset}_train.json"),
        "val": (root / "train2014", root / dataset / f"{mode}_{dataset}_val.json"),
    }

    img_folder, ann_file = PATHS[image_set]
    dataset = ModulatedDetection(
        img_folder,
        ann_file,
        image_set=image_set,
        return_masks=True,
    )
    return dataset


def main():
    # init opts
    args = opts.get_arguments()
    print("Args:", args)

    # create output path
    output_path = './outputs/' + args.outdir
    if not os.path.exists('./outputs/'):
        os.mkdir('./outputs/')
    if not os.path.exists(output_path):
        os.mkdir(output_path)

    # init logger
    if args.train_decoder:
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            filename=os.path.join(output_path, str(args.lr_decoder) + 'decoder_log.txt'),
            filemode='a',
        )
    else:
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            filename=os.path.join(output_path, 'log.txt'),
            filemode='a',
        )

    logger = logging.getLogger()
    console_handler = logging.StreamHandler()
    logger.addHandler(console_handler)

    logger.info('arguments:')
    for arg in vars(args):
        logger.info(f'{arg}: {getattr(args, arg)}')

    # init distribute
    utils.init_distributed_mode(args)
    device = torch.device(args.device)

    # load model
    if args.distributed and dist.get_rank() == 0:
        logger.info("======> load model")
    text_model_name = args.text_encoder
    model = refersam.Model(args, text_model_name, logger).to('cuda')

    ## for clip save dir
    if text_model_name == "ViT-L-14-336px.pt":
        text_model_name = "ViT-L"

    # model = nn.DataParallel(model)
    model.to(device)
    # logger.info(f'model structure: {model}')

    model_without_ddp = model
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True) # for use memory
        model_without_ddp = model.module

    # set param groups
    param_groups = []
    add_to_resizer_param_groups = []
    add_to_mask_decoder_param_groups = []
    add_to_dense_conv_param_groups = []
    add_to_lora_param_groups = []
    add_to_memory_mask_conv_groups = []
    for name, param in model_without_ddp.named_parameters():
        if "resizer" in name:
            param.requires_grad = True
            add_to_resizer_param_groups.append(param)
        elif "dense" in name and "text_encoder" not in name: # for dense_conv and sparse_fc
            param.requires_grad = True
            add_to_dense_conv_param_groups.append(param)
        elif args.train_decoder and 'sam.mask_decoder' in name:
            param.requires_grad = True
            add_to_mask_decoder_param_groups.append(param)
        elif args.train_image_encoder_lora and ('w_a' in name or 'w_b' in name or 'w_a_qkv' in name or 'w_b_qkv' in name):
            param.requires_grad = True
            add_to_lora_param_groups.append(param)        
        elif 'memory_key' in name or 'memory_value' in name:
            param.requires_grad = True
            add_to_memory_mask_conv_groups.append(param)  
        else:
            param.requires_grad = False
    param_groups.append(
        {
            "params": add_to_resizer_param_groups, "lr": args.lr
        }
    )
    if len(add_to_memory_mask_conv_groups) != 0:
        param_groups.append(
            {
                'params': add_to_memory_mask_conv_groups, 'lr': args.lr_memory
            }
        )
    if args.train_decoder:
        param_groups.append(
            {
                'params': add_to_mask_decoder_param_groups, 'lr': args.lr_decoder
            }
        )
    if args.spatial_dynamic_fusion:
        param_groups.append(
            {
                'params': add_to_dense_conv_param_groups, 'lr': args.lr_dense_conv
            }
        )
    if args.train_image_encoder_lora:
        param_groups.append(
            {
                'params': add_to_lora_param_groups, 'lr': args.lr_image_encoder_lora
            }
        )

    # check trained parameters
    logger.info("trained parameters:\n")
    for name, param in model_without_ddp.named_parameters():
        if param.requires_grad:
            logger.info(name)

    if args.distributed and dist.get_rank() == 0:
        n_parameters = sum(p.numel() for p in model_without_ddp.parameters() if p.requires_grad)
        logger.info('number of params:', n_parameters)

        logger.info("======> train config")
    # load train dataset
    data_file = "all"
    dataset_names = ["refcoco", "refcoco+", "refcocog"]
    train_dataset = torch.utils.data.ConcatDataset(
        [build_refexp(name, image_set="train") for name in dataset_names]
    )

    if args.distributed and dist.get_rank() == 0:
        logger.info("\nTrain dataset sample number: ", len(train_dataset))
        logger.info("\n")

    optimizer = torch.optim.AdamW(param_groups, lr=args.lr)

    if args.distributed:
        if args.cache_mode:
            sampler_train = samplers.NodeDistributedSampler(train_dataset)
        else:
            sampler_train = samplers.DistributedSampler(train_dataset)
    else:
        sampler_train = torch.utils.data.RandomSampler(train_dataset) #02
    
    batch_sampler_train = torch.utils.data.BatchSampler(
        sampler_train, args.batch_size, drop_last=True)

    train_dataloader = DataLoader(train_dataset, batch_sampler=batch_sampler_train, collate_fn=utils.collate_fn, num_workers=args.num_workers)

    # load valid dataset
    Val_all = namedtuple(typename="val_data", field_names=["dataset_name", "dataloader", "base_ds", "evaluator_list"])
    val_tuples = []

    for name in dataset_names:
        val_dataset = build_refexp(name, image_set="val")
        sampler_val = (
            samplers.DistributedSampler(val_dataset, shuffle=False) if args.distributed else torch.utils.data.SequentialSampler(val_dataset)
        )
        data_loader_val = DataLoader(
            val_dataset,
            args.batch_size,
            sampler=sampler_val,
            drop_last=False,
            collate_fn=utils.collate_fn,
            num_workers=args.num_workers
        )
        base_ds = get_coco_api_from_dataset(val_dataset)
        val_tuples.append(Val_all(dataset_name=name, dataloader=data_loader_val, base_ds=base_ds, evaluator_list=None))

    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.train_epoch)
    if args.distributed and dist.get_rank() == 0:
        logger.info("======> staring training")
    for train_idx in tqdm(range(args.train_epoch)):
        if args.distributed:
            sampler_train.set_epoch(train_idx)

        # model.module.resizer.train()
        model.train()
        one_epoch_loss = 0
        sampler_num = 0

        for idx, (img, mask, caption, target) in enumerate(train_dataloader):

            img = img.to(device)
            mask = mask.to(device)
            train_mask_list = mask

            high_res_masks_list = model(img, caption, target)

            one_iter_sampler = 0
            one_loss = 0
            if isinstance(high_res_masks_list, list):
                high_res_masks_list = torch.cat(high_res_masks_list, dim=0)

            for high_res_masks, train_mask in zip(high_res_masks_list, train_mask_list):
                high_res_masks = high_res_masks.flatten(1)

                gt_mask = train_mask
                if isinstance(gt_mask, torch.Tensor):
                    gt_mask = gt_mask.cpu().float().unsqueeze(0).flatten(1).cuda()
                else:
                    gt_mask = torch.Tensor(gt_mask.cpu()).float().unsqueeze(0).flatten(1).cuda()

                dice_loss = loss.calculate_dice_loss(high_res_masks, gt_mask)  
                focal_loss = loss.calculate_sigmoid_focal_loss(high_res_masks, gt_mask)  
                # point_loss = loss_fn(after_text_sentence_feature, target_feat)
                one_loss += dice_loss + focal_loss
                one_iter_sampler += 1

            one_loss /= one_iter_sampler
            sampler_num += 1
            one_epoch_loss += one_loss.item()
            optimizer.zero_grad()
            one_loss.backward()
            optimizer.step()

            # check weight
            # logger.info("weight of lora layer:{}".format(model_without_ddp.sam.image_encoder.blocks[-1].attn.proj.w_b.weight))

            if (args.distributed and dist.get_rank() == 0) or not args.distributed:
                if idx % 10 == 0:
                    logger.info("index:{}, total_sum:{}, loss:{}, lr:{}".format(idx, len(train_dataloader), one_loss, args.lr))

        # scheduler.step()

        one_epoch_loss /= sampler_num
        # logger.info("loss_sum:{}".format(loss))
        if (args.distributed and dist.get_rank() == 0) or not args.distributed:
            if train_idx % args.log_epoch == 0:
                logger.info('Train Epoch: {:} / {:}'.format(train_idx, args.train_epoch))
                # current_lr = scheduler.get_last_lr()[0]
                # logger.info('LR: {:.6f}, Loss: {:.4f}'.format(current_lr, one_epoch_loss))
                logger.info('LR: {:.6f}, Loss: {:.4f}'.format(args.lr, one_epoch_loss))

        if (args.distributed and dist.get_rank() == 0) or not args.distributed:
            # save last model
            torch.save(model_without_ddp.resizer.state_dict(),
                       os.path.join(output_path, 'resizer_' + text_model_name + '_last.pth'))
            if args.train_decoder:
                torch.save(model_without_ddp.sam.mask_decoder.state_dict(),
                           os.path.join(output_path, 'mask_decoder_' + text_model_name + '_last.pth'))
            if args.spatial_dynamic_fusion:
                torch.save(model_without_ddp.dense_conv.state_dict(),
                           os.path.join(output_path, 'dense_conv_' + text_model_name + '_last.pth'))
            if args.sparse_attention:
                torch.save(model_without_ddp.sparse_fc.state_dict(),
                           os.path.join(output_path, 'sparse_fc' + text_model_name + '_last.pth'))
            if args.train_image_encoder_lora:
                torch.save(model_without_ddp.sam.image_encoder.blocks.state_dict(),
                           os.path.join(output_path, 'image_encoder_blocks_' + text_model_name + '_last.pth'))
            if args.word_memory:
                torch.save(model_without_ddp.memory_key.state_dict(),
                           os.path.join(output_path, 'memory_key_' + text_model_name + '_last.pth'))
                torch.save(model_without_ddp.memory_value.state_dict(),
                           os.path.join(output_path, 'memory_value_' + text_model_name + '_last.pth'))
 
            # each N epoch save one model
            if train_idx % 1 == 0:
                torch.save(model_without_ddp.resizer.state_dict(), os.path.join(output_path, 'resizer_' + text_model_name + "_" + str(train_idx) + '.pth'))
                if args.train_decoder:
                    torch.save(model_without_ddp.sam.mask_decoder.state_dict(), os.path.join(output_path, 'mask_decoder_' + text_model_name + "_" + str(train_idx) + '.pth'))
                if args.spatial_dynamic_fusion:
                    torch.save(model_without_ddp.dense_conv.state_dict(), os.path.join(output_path, 'dense_conv_' + text_model_name + "_" + str(train_idx) + '.pth'))
                if args.sparse_attention:
                    torch.save(model_without_ddp.sparse_fc.state_dict(), os.path.join(output_path, 'sparse_fc' + text_model_name + "_" + str(train_idx) + '.pth'))
                if args.train_image_encoder_lora:
                    torch.save(model_without_ddp.sam.image_encoder.blocks.state_dict(),
                               os.path.join(output_path, 'image_encoder_blocks_' + text_model_name + "_" + str(train_idx) + '.pth'))
                if args.word_memory:
                    torch.save(model_without_ddp.memory_key.state_dict(),
                            os.path.join(output_path, 'memory_key_' + text_model_name + "_" + str(train_idx) + '.pth'))
                    torch.save(model_without_ddp.memory_value.state_dict(),
                            os.path.join(output_path, 'memory_value_' + text_model_name + "_" + str(train_idx) + '.pth'))


        # each epoch eval
        test_stats = {}
        for i, item in enumerate(val_tuples):
            evaluator_list = build_evaluator_list(item.base_ds, item.dataset_name)
            item = item._replace(evaluator_list=evaluator_list)
            logger.info(f"\n Evaluating {item.dataset_name}")

            model.eval()
            one_epoch_loss = 0
            sampler_num = 0

            for idx, (img, mask, caption, target) in enumerate(item.dataloader):
                img = img.to(device)
                mask = mask.to(device)
                train_mask_list = mask

                with torch.no_grad():
                    # high_res_masks_list = model(target)
                    high_res_masks_list = model(img, caption, target)

                one_iter_sampler = 0
                one_loss = 0
                if isinstance(high_res_masks_list, list):
                    high_res_masks_list = torch.cat(high_res_masks_list, dim=0)

                for high_res_masks, train_mask in zip(high_res_masks_list, train_mask_list):
                    high_res_masks = high_res_masks.flatten(1)

                    gt_mask = train_mask
                    if isinstance(gt_mask, torch.Tensor):
                        gt_mask = gt_mask.cpu().float().unsqueeze(0).flatten(1).cuda()
                    else:
                        gt_mask = torch.Tensor(gt_mask.cpu()).float().unsqueeze(0).flatten(1).cuda()

                    dice_loss = loss.calculate_dice_loss(high_res_masks, gt_mask)  # 用于图像分割中的目标边界分割 < 0.1较好
                    focal_loss = loss.calculate_sigmoid_focal_loss(high_res_masks, gt_mask)  # 用于像素级别分类
                    one_loss += dice_loss + focal_loss
                    one_iter_sampler += 1

                one_loss /= one_iter_sampler
                sampler_num += 1
                one_epoch_loss += one_loss.item()                

                if idx % 100 == 0:
                    logger.info("eval index:{}, total_sum:{}, loss:{}".format(idx, len(item.dataloader), one_loss))

            one_epoch_loss /= sampler_num

            if train_idx % args.log_epoch == 0:
                logger.info('Train Epoch: {:} / {:}'.format(train_idx, args.train_epoch))
                logger.info('eval dataset {} Loss: {:.4f}'.format(item.dataset_name, one_epoch_loss))

            # if not args.distributed or dist.get_rank() == 0:
                # logger.info("proj w_b weight:".format(model.sam.image_encoder.blocks[-1].attn.proj.w_b.weight))
                # logger.info("qkv w_b_qkv weight:".format(model.sam.image_encoder.blocks[-1].attn.qkv.w_b_qkv.weight))

    if not args.distributed or dist.get_rank() == 0:
        torch.save(model_without_ddp.resizer.state_dict(),
                   os.path.join(output_path, 'resizer_' + text_model_name + "_" + str(args.train_epoch) + '.pth'))
        if args.train_decoder:
            torch.save(model_without_ddp.sam.mask_decoder.state_dict(),
                       os.path.join(output_path, 'mask_decoder_' + text_model_name + "_" + str(args.train_epoch) + '.pth'))
        if args.spatial_dynamic_fusion:
            torch.save(model_without_ddp.dense_conv.state_dict(),
                       os.path.join(output_path, "dense_conv_" + text_model_name + "_" + str(args.train_epoch) + '.pth'))
        if args.sparse_attention:
            torch.save(model_without_ddp.sparse_fc.state_dict(), 
                       os.path.join(output_path, 'sparse_fc' + text_model_name + "_" + str(args.train_epoch) + '.pth'))
        if args.train_image_encoder_lora:
            torch.save(model_without_ddp.sam.image_encoder.blocks.state_dict(),
                       os.path.join(output_path, 'image_encoder_blocks_' + text_model_name + "_" + str(args.train_epoch) + '.pth'))
        if args.word_memory:
            torch.save(model_without_ddp.memory_key.state_dict(),
                    os.path.join(output_path, 'memory_key_' + text_model_name + "_" + str(args.train_epoch) + '.pth'))
            torch.save(model_without_ddp.memory_value.state_dict(),
                    os.path.join(output_path, 'memory_value_' + text_model_name + "_" + str(args.train_epoch) + '.pth'))


def get_coco_api_from_dataset(dataset):
    for _ in range(10):
        # if isinstance(dataset, torchvision.datasets.CocoDetection):
        #     break
        if isinstance(dataset, torch.utils.data.Subset):
            dataset = dataset.dataset
    if isinstance(dataset, torchvision.datasets.CocoDetection):
        return dataset.coco


# build evaluator list for dataset_val
def build_evaluator_list(base_ds, dataset_name):
    """Helper function to build the list of evaluators for a given dataset"""
    evaluator_list = []
    iou_types = ["bbox"]
    iou_types.append("segm")

    evaluator_list.append(CocoEvaluator(base_ds, tuple(iou_types), useCats=False))
    # TODO: currently ont support RefExpEvaluator (memory error)
    return evaluator_list


if __name__ == "__main__":
    main()