Release training script

Former-commit-id: 96ec3cdf6a4f6880ac274ddde55537570788ebbb

model/LISA.py

@@ -3,14 +3,18 @@ from typing import List
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from peft import (LoraConfig, get_peft_model)
+from peft import LoraConfig, get_peft_model
 from transformers import BitsAndBytesConfig, CLIPVisionModel
 
 from transformers import CLIPVisionModel, BitsAndBytesConfig
 from .llava.model.llava import LlavaLlamaForCausalLM
 from .segment_anything import build_sam_vit_h
-from utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
-                    DEFAULT_IMAGE_PATCH_TOKEN)
+from utils.utils import (
+    DEFAULT_IM_END_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IMAGE_PATCH_TOKEN,
+)
+
 
 def dice_loss(
     inputs: torch.Tensor,
@@ -219,7 +223,9 @@ class LISA(nn.Module):
         self.lm.resize_token_embeddings(len(tokenizer))
 
         for n, p in self.lm.named_parameters():
-            if any([x in n for x in ["lm_head", "embed_tokens"]]) and p.shape[0] == len(tokenizer):
+            if any([x in n for x in ["lm_head", "embed_tokens"]]) and p.shape[0] == len(
+                tokenizer
+            ):
                 p.requires_grad = True
 
         # SAM

model/llava/eval/model_vqa.py

@@ -11,9 +11,14 @@ from llava.conversation import conv_templates
 from llava.utils import disable_torch_init
 from PIL import Image
 from tqdm import tqdm
-from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
-                          CLIPImageProcessor, CLIPVisionModel,
-                          StoppingCriteria)
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    CLIPImageProcessor,
+    CLIPVisionModel,
+    StoppingCriteria,
+)
 
 
 def split_list(lst, n):

model/llava/eval/model_vqa_science.py

@@ -11,9 +11,14 @@ from llava.conversation import conv_templates
 from llava.utils import disable_torch_init
 from PIL import Image
 from tqdm import tqdm
-from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
-                          CLIPImageProcessor, CLIPVisionModel,
-                          StoppingCriteria)
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    CLIPImageProcessor,
+    CLIPVisionModel,
+    StoppingCriteria,
+)
 
 
 def split_list(lst, n):

model/llava/eval/run_llava.py

@@ -9,9 +9,13 @@ from llava.model import *
 from llava.model.utils import KeywordsStoppingCriteria
 from llava.utils import disable_torch_init
 from PIL import Image
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          CLIPImageProcessor, CLIPVisionModel,
-                          StoppingCriteria)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    CLIPImageProcessor,
+    CLIPVisionModel,
+    StoppingCriteria,
+)
 
 DEFAULT_IMAGE_TOKEN = "<image>"
 DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"

model/llava/eval/run_llava_batch.py

@@ -13,9 +13,13 @@ from llava.model import *
 from llava.model.utils import KeywordsStoppingCriteria
 from llava.utils import disable_torch_init
 from PIL import Image
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          CLIPImageProcessor, CLIPVisionModel,
-                          StoppingCriteria)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    CLIPImageProcessor,
+    CLIPVisionModel,
+    StoppingCriteria,
+)
 
 DEFAULT_IMAGE_TOKEN = "<image>"
 DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"

model/llava/eval/run_llava_batch_v2.py

@@ -13,9 +13,13 @@ from llava.model import *
 from llava.model.utils import KeywordsStoppingCriteria
 from llava.utils import disable_torch_init
 from PIL import Image
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          CLIPImageProcessor, CLIPVisionModel,
-                          StoppingCriteria)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    CLIPImageProcessor,
+    CLIPVisionModel,
+    StoppingCriteria,
+)
 
 DEFAULT_IMAGE_TOKEN = "<image>"
 DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"

model/llava/eval/run_llava_batch_v3.py

@@ -13,9 +13,13 @@ from llava.model import *
 from llava.model.utils import KeywordsStoppingCriteria
 from llava.utils import disable_torch_init
 from PIL import Image
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          CLIPImageProcessor, CLIPVisionModel,
-                          StoppingCriteria)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    CLIPImageProcessor,
+    CLIPVisionModel,
+    StoppingCriteria,
+)
 
 DEFAULT_IMAGE_TOKEN = "<image>"
 DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"

model/llava/model/llava.py

@@ -19,11 +19,19 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn import CrossEntropyLoss
-from transformers import (AutoConfig, AutoModelForCausalLM, CLIPImageProcessor,
-                          CLIPVisionModel, LlamaConfig, LlamaForCausalLM,
-                          LlamaModel)
-from transformers.modeling_outputs import (BaseModelOutputWithPast,
-                                           CausalLMOutputWithPast)
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    CLIPImageProcessor,
+    CLIPVisionModel,
+    LlamaConfig,
+    LlamaForCausalLM,
+    LlamaModel,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
 
 DEFAULT_IMAGE_TOKEN = "<image>"
 DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"

model/llava/model/llava_mpt.py

@@ -21,10 +21,16 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn import CrossEntropyLoss
-from transformers import (AutoConfig, AutoModelForCausalLM, CLIPImageProcessor,
-                          CLIPVisionModel)
-from transformers.modeling_outputs import (BaseModelOutputWithPast,
-                                           CausalLMOutputWithPast)
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    CLIPImageProcessor,
+    CLIPVisionModel,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
 
 from .mpt.modeling_mpt import MPTConfig, MPTForCausalLM, MPTModel
 

model/llava/model/mpt/adapt_tokenizer.py

@@ -1,7 +1,6 @@
 from typing import Union
 
-from transformers import (AutoTokenizer, PreTrainedTokenizer,
-                          PreTrainedTokenizerFast)
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
 
 Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
 NUM_SENTINEL_TOKENS: int = 100

model/llava/model/mpt/hf_prefixlm_converter.py

@@ -13,22 +13,26 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 from transformers.models.bloom.modeling_bloom import (
-    BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel,
-    CausalLMOutputWithCrossAttentions, CrossEntropyLoss)
-from transformers.models.bloom.modeling_bloom import \
-    _expand_mask as _expand_mask_bloom
-from transformers.models.bloom.modeling_bloom import \
-    _make_causal_mask as _make_causal_mask_bloom
+    BaseModelOutputWithPastAndCrossAttentions,
+    BloomForCausalLM,
+    BloomModel,
+    CausalLMOutputWithCrossAttentions,
+    CrossEntropyLoss,
+)
+from transformers.models.bloom.modeling_bloom import _expand_mask as _expand_mask_bloom
+from transformers.models.bloom.modeling_bloom import (
+    _make_causal_mask as _make_causal_mask_bloom,
+)
 from transformers.models.bloom.modeling_bloom import logging
 from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
 from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
 from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
 from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
 from transformers.models.opt.modeling_opt import OPTForCausalLM
-from transformers.models.opt.modeling_opt import \
-    _expand_mask as _expand_mask_opt
-from transformers.models.opt.modeling_opt import \
-    _make_causal_mask as _make_causal_mask_opt
+from transformers.models.opt.modeling_opt import _expand_mask as _expand_mask_opt
+from transformers.models.opt.modeling_opt import (
+    _make_causal_mask as _make_causal_mask_opt,
+)
 
 logger = logging.get_logger(__name__)
 _SUPPORTED_GPT_MODELS = (

model/llava/model/mpt/modeling_mpt.py

@@ -9,17 +9,20 @@ from typing import List, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import (PreTrainedModel, PreTrainedTokenizer,
-                          PreTrainedTokenizerFast)
-from transformers.modeling_outputs import (BaseModelOutputWithPast,
-                                           CausalLMOutputWithPast)
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
 
 from .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising
 from .attention import attn_bias_shape, build_attn_bias
 from .blocks import MPTBlock
 from .configuration_mpt import MPTConfig
-from .hf_prefixlm_converter import (add_bidirectional_mask_if_missing,
-                                    convert_hf_causal_lm_to_prefix_lm)
+from .hf_prefixlm_converter import (
+    add_bidirectional_mask_if_missing,
+    convert_hf_causal_lm_to_prefix_lm,
+)
 from .meta_init_context import init_empty_weights
 from .norm import NORM_CLASS_REGISTRY
 from .param_init_fns import MODEL_INIT_REGISTRY, generic_param_init_fn_

model/llava/serve/gradio_web_server.py

@@ -9,12 +9,15 @@ from collections import defaultdict
 import gradio as gr
 import requests
 from llava.constants import LOGDIR
-from llava.conversation import (SeparatorStyle, conv_templates,
-                                default_conversation)
+from llava.conversation import SeparatorStyle, conv_templates, default_conversation
 from llava.serve.gradio_css import code_highlight_css
 from llava.serve.gradio_patch import Chatbot as grChatbot
-from llava.utils import (build_logger, moderation_msg, server_error_msg,
-                         violates_moderation)
+from llava.utils import (
+    build_logger,
+    moderation_msg,
+    server_error_msg,
+    violates_moderation,
+)
 
 logger = build_logger("gradio_web_server", "gradio_web_server.log")
 

model/llava/train/train.py

@@ -715,8 +715,9 @@ def train():
                     "[WARNING] As of 4/30/23, this feature requires PyTorch-nightly build.  See here for details: https://github.com/haotian-liu/LLaVA#experimental-use-fsdp-to-save-memory-in-pretraining"
                 )
 
-                from torch.distributed.fsdp.fully_sharded_data_parallel import \
-                    FullyShardedDataParallel as FSDP
+                from torch.distributed.fsdp.fully_sharded_data_parallel import (
+                    FullyShardedDataParallel as FSDP,
+                )
 
                 def patch_FSDP_use_orig_params(func):
                     def wrap_func(*args, **kwargs):

model/llava/train/train_mem.py

@@ -3,8 +3,7 @@
 # Make it more memory efficient by monkey patching the LLaMA model with FlashAttn.
 
 # Need to call this before importing transformers.
-from llava.train.llama_flash_attn_monkey_patch import \
-    replace_llama_attn_with_flash_attn
+from llava.train.llama_flash_attn_monkey_patch import replace_llama_attn_with_flash_attn
 
 replace_llama_attn_with_flash_attn()
 

model/segment_anything/__init__.py

@@ -5,6 +5,11 @@
 # LICENSE file in the root directory of this source tree.
 
 from .automatic_mask_generator import SamAutomaticMaskGenerator
-from .build_sam import (build_sam, build_sam_vit_b, build_sam_vit_h,
-                        build_sam_vit_l, sam_model_registry)
+from .build_sam import (
+    build_sam,
+    build_sam_vit_b,
+    build_sam_vit_h,
+    build_sam_vit_l,
+    sam_model_registry,
+)
 from .predictor import SamPredictor

model/segment_anything/automatic_mask_generator.py

@@ -12,13 +12,24 @@ from torchvision.ops.boxes import batched_nms, box_area  # type: ignore
 
 from .modeling import Sam
 from .predictor import SamPredictor
-from .utils.amg import (MaskData, area_from_rle, batch_iterator,
-                        batched_mask_to_box, box_xyxy_to_xywh,
-                        build_all_layer_point_grids, calculate_stability_score,
-                        coco_encode_rle, generate_crop_boxes,
-                        is_box_near_crop_edge, mask_to_rle_pytorch,
-                        remove_small_regions, rle_to_mask, uncrop_boxes_xyxy,
-                        uncrop_masks, uncrop_points)
+from .utils.amg import (
+    MaskData,
+    area_from_rle,
+    batch_iterator,
+    batched_mask_to_box,
+    box_xyxy_to_xywh,
+    build_all_layer_point_grids,
+    calculate_stability_score,
+    coco_encode_rle,
+    generate_crop_boxes,
+    is_box_near_crop_edge,
+    mask_to_rle_pytorch,
+    remove_small_regions,
+    rle_to_mask,
+    uncrop_boxes_xyxy,
+    uncrop_masks,
+    uncrop_points,
+)
 
 
 class SamAutomaticMaskGenerator:
@@ -104,8 +115,7 @@ class SamAutomaticMaskGenerator:
             "coco_rle",
         ], f"Unknown output_mode {output_mode}."
         if output_mode == "coco_rle":
-            from pycocotools import \
-                mask as mask_utils  # type: ignore # noqa: F401
+            from pycocotools import mask as mask_utils  # type: ignore # noqa: F401
 
         if min_mask_region_area > 0:
             import cv2  # type: ignore # noqa: F401

model/segment_anything/build_sam.py

@@ -8,8 +8,13 @@ from functools import partial
 
 import torch
 
-from .modeling import (ImageEncoderViT, MaskDecoder, PromptEncoder, Sam,
-                       TwoWayTransformer)
+from .modeling import (
+    ImageEncoderViT,
+    MaskDecoder,
+    PromptEncoder,
+    Sam,
+    TwoWayTransformer,
+)
 
 
 def build_sam_vit_h(checkpoint=None):

train_ds.py

@@ -14,8 +14,13 @@ from torch.utils.tensorboard import SummaryWriter
 
 from model.LISA import LISA
 from utils.dataset import HybridDataset, ValDataset, collate_fn
-from utils.utils import (AverageMeter, ProgressMeter, Summary, dict_to_cuda,
-                         intersectionAndUnionGPU)
+from utils.utils import (
+    AverageMeter,
+    ProgressMeter,
+    Summary,
+    dict_to_cuda,
+    intersectionAndUnionGPU,
+)
 
 
 def parse_args(args):
@@ -54,9 +59,7 @@ def parse_args(args):
     )
     parser.add_argument("--vqa_data", default="llava_instruct_150k", type=str)
     parser.add_argument("--reason_seg_data", default="ReasonSeg|train", type=str)
-    parser.add_argument(
-        "--val_dataset", default="ReasonSeg|val", type=str
-    )
+    parser.add_argument("--val_dataset", default="ReasonSeg|val", type=str)
     parser.add_argument("--dataset_dir", default="./dataset", type=str)
     parser.add_argument("--log_base_dir", default="./runs", type=str)
     parser.add_argument("--exp_name", default="lisa", type=str)
@@ -87,7 +90,9 @@ def parse_args(args):
     parser.add_argument("--exclude_val", action="store_true", default=False)
     parser.add_argument("--no_eval", action="store_true", default=False)
     parser.add_argument("--eval_only", action="store_true", default=False)
-    parser.add_argument("--vision_pretrained", default="PATH TO SAM ViT-H Pre-trained Wegiht", type=str)
+    parser.add_argument(
+        "--vision_pretrained", default="PATH TO SAM ViT-H Pre-trained Wegiht", type=str
+    )
     parser.add_argument("--weight", default="", type=str)
     parser.add_argument("--print_freq", default=1, type=int)
     parser.add_argument("--start_epoch", default=0, type=int)
@@ -133,7 +138,7 @@ def main(args):
     )
 
     if args.weight:
-        state_dict = torch.load(args.weight, map_location='cpu')
+        state_dict = torch.load(args.weight, map_location="cpu")
         model.load_state_dict(state_dict, strict=True)
 
     world_size = torch.cuda.device_count()
@@ -142,7 +147,10 @@ def main(args):
         args.dataset_dir,
         tokenizer,
         args.vision_tower,
-        samples_per_epoch=args.batch_size * args.grad_accumulation_steps * args.steps_per_epoch * world_size,
+        samples_per_epoch=args.batch_size
+        * args.grad_accumulation_steps
+        * args.steps_per_epoch
+        * world_size,
         precision=args.precision,
         image_size=args.image_size,
         num_classes_per_sample=args.num_classes_per_sample,
@@ -163,7 +171,9 @@ def main(args):
             args.val_dataset,
             args.image_size,
         )
-        print(f"Training with {len(train_dataset)} examples and validating with {len(val_dataset)} examples.")
+        print(
+            f"Training with {len(train_dataset)} examples and validating with {len(val_dataset)} examples."
+        )
     else:
         val_dataset = None
         print(f"Training with {len(train_dataset)} examples.")
@@ -215,7 +225,9 @@ def main(args):
 
     if val_dataset is not None:
         assert args.val_batch_size == 1
-        val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False, drop_last=False)
+        val_sampler = torch.utils.data.distributed.DistributedSampler(
+            val_dataset, shuffle=False, drop_last=False
+        )
         val_loader = torch.utils.data.DataLoader(
             val_dataset,
             batch_size=args.val_batch_size,
@@ -230,13 +242,10 @@ def main(args):
     best_score, cur_ciou = 0.0, 0.0
 
     if args.eval_only:
-        giou, ciou = validate(
-            val_loader, model_engine, 0, writer, args
-        )
+        giou, ciou = validate(val_loader, model_engine, 0, writer, args)
         exit()
 
     for epoch in range(args.start_epoch, args.epochs):
-
         # train for one epoch
         train_iter = train(
             train_loader,
@@ -249,9 +258,7 @@ def main(args):
         )
 
         if args.no_eval == False:
-            giou, ciou = validate(
-                val_loader, model_engine, epoch, writer, args
-            )
+            giou, ciou = validate(val_loader, model_engine, epoch, writer, args)
             is_best = giou > best_score
             best_score = max(giou, best_score)
             cur_ciou = ciou if is_best else cur_ciou

utils/dataset.py

@@ -17,8 +17,12 @@ from .reason_seg_dataset import ReasonSegDataset
 from .refer import REFER
 from .refer_seg_dataset import ReferSegDataset
 from .sem_seg_dataset import SemSegDataset
-from .utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
-                    DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IMAGE_TOKEN)
+from .utils import (
+    DEFAULT_IM_END_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IMAGE_PATCH_TOKEN,
+    DEFAULT_IMAGE_TOKEN,
+)
 from .vqa_dataset import VQADataset
 
 
@@ -67,7 +71,7 @@ def collate_fn(batch, tokenizer=None):
         max_length=tokenizer.model_max_length,
         truncation=True,
     )
-        
+
     input_ids = tokenize_data.input_ids
     attention_masks = tokenize_data.attention_mask
 
@@ -261,7 +265,7 @@ class ValDataset(torch.utils.data.Dataset):
                 os.path.join(self.base_image_dir, "reason_seg", ds, split, "*.jpg")
             )
             self.images = images
-            self.data_type = 'reason_seg'
+            self.data_type = "reason_seg"
         elif len(splits) == 3:
             ds, splitBy, split = splits
             refer_api = REFER(self.base_image_dir, ds, splitBy)
@@ -294,7 +298,7 @@ class ValDataset(torch.utils.data.Dataset):
                 ]
             refer_seg_ds["img2refs"] = img2refs
             self.refer_seg_ds = refer_seg_ds
-            self.data_type = 'refer_seg'
+            self.data_type = "refer_seg"
 
         self.ds = ds
         self.image_size = image_size
@@ -303,7 +307,7 @@ class ValDataset(torch.utils.data.Dataset):
         self.clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower)
 
     def __len__(self):
-        if self.data_type == 'refer_seg':
+        if self.data_type == "refer_seg":
             return len(self.refer_seg_ds["images"])
         else:
             return len(self.images)
@@ -321,7 +325,7 @@ class ValDataset(torch.utils.data.Dataset):
         return x
 
     def __getitem__(self, idx):
-        if self.data_type == 'refer_seg':
+        if self.data_type == "refer_seg":
             refer_seg_ds = self.refer_seg_ds
             images = refer_seg_ds["images"]
             annotations = refer_seg_ds["annotations"]
@@ -406,7 +410,7 @@ class ValDataset(torch.utils.data.Dataset):
 
         images = self.preprocess(torch.from_numpy(images).permute(2, 0, 1).contiguous())
 
-        if self.data_type == 'refer_seg':
+        if self.data_type == "refer_seg":
             masks = []
             for i, ann_id in enumerate(sampled_ann_ids):
                 ann = annotations[ann_id]

utils/reason_seg_dataset.py

@@ -13,10 +13,16 @@ from model.segment_anything.utils.transforms import ResizeLongestSide
 
 from .conversation import get_default_conv_template
 from .data_processing import get_mask_from_json
-from .utils import (ANSWER_LIST, DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
-                    DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IMAGE_TOKEN,
-                    EXPLANATORY_QUESTION_LIST, LONG_QUESTION_LIST,
-                    SHORT_QUESTION_LIST)
+from .utils import (
+    ANSWER_LIST,
+    DEFAULT_IM_END_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IMAGE_PATCH_TOKEN,
+    DEFAULT_IMAGE_TOKEN,
+    EXPLANATORY_QUESTION_LIST,
+    LONG_QUESTION_LIST,
+    SHORT_QUESTION_LIST,
+)
 
 
 class ReasonSegDataset(torch.utils.data.Dataset):
@@ -72,7 +78,13 @@ class ReasonSegDataset(torch.utils.data.Dataset):
             self.explanatory_question_list = EXPLANATORY_QUESTION_LIST
             self.img_to_explanation = {}
             with open(
-                os.path.join(base_image_dir, "reason_seg", reason_seg_data, "explanatory", "train.json")
+                os.path.join(
+                    base_image_dir,
+                    "reason_seg",
+                    reason_seg_data,
+                    "explanatory",
+                    "train.json",
+                )
             ) as f:
                 items = json.load(f)
             for item in items:
@@ -131,9 +143,7 @@ class ReasonSegDataset(torch.utils.data.Dataset):
         ]
 
         image_name = image_path.split("/")[-1]
-        if (
-            self.explanatory != -1 and image_name in self.img_to_explanation
-        ): 
+        if self.explanatory != -1 and image_name in self.img_to_explanation:
             if random.random() < self.explanatory:
                 choice = 2
             else:
@@ -200,7 +210,11 @@ class ReasonSegDataset(torch.utils.data.Dataset):
         images = self.preprocess(torch.from_numpy(images).permute(2, 0, 1).contiguous())
 
         image_name = image_path.split("/")[-1]
-        if self.explanatory != -1 and image_name in self.img_to_explanation and choice == 2:
+        if (
+            self.explanatory != -1
+            and image_name in self.img_to_explanation
+            and choice == 2
+        ):
             masks = torch.rand(0, *ori_size)
             label = torch.ones(ori_size) * self.ignore_label
         else:

utils/refer_seg_dataset.py

@@ -12,9 +12,14 @@ from model.segment_anything.utils.transforms import ResizeLongestSide
 
 from .conversation import get_default_conv_template
 from .refer import REFER
-from .utils import (ANSWER_LIST, DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
-                    DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IMAGE_TOKEN,
-                    SHORT_QUESTION_LIST)
+from .utils import (
+    ANSWER_LIST,
+    DEFAULT_IM_END_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IMAGE_PATCH_TOKEN,
+    DEFAULT_IMAGE_TOKEN,
+    SHORT_QUESTION_LIST,
+)
 
 
 class ReferSegDataset(torch.utils.data.Dataset):

utils/sem_seg_dataset.py

@@ -14,9 +14,15 @@ from transformers import CLIPImageProcessor
 from model.segment_anything.utils.transforms import ResizeLongestSide
 
 from .conversation import get_default_conv_template
-from .utils import (ANSWER_LIST, DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
-                    DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IMAGE_TOKEN,
-                    SHORT_QUESTION_LIST)
+from .utils import (
+    ANSWER_LIST,
+    DEFAULT_IM_END_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IMAGE_PATCH_TOKEN,
+    DEFAULT_IMAGE_TOKEN,
+    SHORT_QUESTION_LIST,
+)
+
 
 def init_mapillary(base_image_dir):
     mapillary_data_root = os.path.join(base_image_dir, "mapillary")

utils/utils.py

@@ -12,8 +12,12 @@ DEFAULT_IM_END_TOKEN = "<im_end>"
 SHORT_QUESTION_LIST = [
     DEFAULT_IMAGE_TOKEN + " " + "Can you segment the {class_name} in this image?",
     DEFAULT_IMAGE_TOKEN + " " + "Please segment the {class_name} in this image.",
-    DEFAULT_IMAGE_TOKEN + " " + "What is {class_name} in this image? Please respond with segmentation mask.",
-    DEFAULT_IMAGE_TOKEN + " " + "What is {class_name} in this image? Please output segmentation mask.",
+    DEFAULT_IMAGE_TOKEN
+    + " "
+    + "What is {class_name} in this image? Please respond with segmentation mask.",
+    DEFAULT_IMAGE_TOKEN
+    + " "
+    + "What is {class_name} in this image? Please output segmentation mask.",
 ]
 
 LONG_QUESTION_LIST = [
@@ -121,6 +125,7 @@ def intersectionAndUnionGPU(output, target, K, ignore_index=255):
     area_union = area_output + area_target - area_intersection
     return area_intersection, area_union, area_target
 
+
 class ProgressMeter(object):
     def __init__(self, num_batches, meters, prefix=""):
         self.batch_fmtstr = self._get_batch_fmtstr(num_batches)

utils/vqa_dataset.py

@@ -10,8 +10,13 @@ from transformers import CLIPImageProcessor
 from model.segment_anything.utils.transforms import ResizeLongestSide
 
 from .conversation import get_default_conv_template
-from .utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
-                    DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IMAGE_TOKEN)
+from .utils import (
+    DEFAULT_IM_END_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IMAGE_PATCH_TOKEN,
+    DEFAULT_IMAGE_TOKEN,
+)
+
 
 class VQADataset(torch.utils.data.Dataset):
     pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
@@ -49,7 +54,7 @@ class VQADataset(torch.utils.data.Dataset):
         self.vqa_data = vqa_data
 
         print("vqa_data: ", len(self.vqa_data))
-        
+
     def __len__(self):
         return self.samples_per_epoch
 
@@ -72,7 +77,11 @@ class VQADataset(torch.utils.data.Dataset):
         img = cv2.imread(image_path)
         images = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         ori_size = images.shape[:2]
-        images_clip = self.clip_image_processor.preprocess(images, return_tensors="pt")["pixel_values"][0]  # preprocess images for clip
+        images_clip = self.clip_image_processor.preprocess(images, return_tensors="pt")[
+            "pixel_values"
+        ][
+            0
+        ]  # preprocess images for clip
         image_token_len = (images_clip.shape[1] // 14) * (
             images_clip.shape[2] // 14
         )  # FIXME: 14 is hardcoded patch size