+
+ Image2Text
+{gen_text}
+
+ 
+
+ Text2Image
+
+ {gen_text}
+
+ 
Understanding Image with Text'
+
+# Create Gradio interface
+interface = gr.Interface(
+ fn=lambda image: process_image(image, processor), # Pass the processor object using a lambda function
+ inputs=image_input,
+ outputs=gr.outputs.HTML(),
+ title=title_with_logo,
+ description="""
+ This code support image to text transformation. Then the generated text can do retrieval, question answering et al to conduct zero-shot.
+ """
+)
+
+# Launch the interface
+interface.launch()
\ No newline at end of file
diff --git a/main_gradio.py b/main_gradio.py
new file mode 100644
index 0000000000000000000000000000000000000000..d567ef65df74449cbd5f917105da5c40bb9cca37
--- /dev/null
+++ b/main_gradio.py
@@ -0,0 +1,84 @@
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+import base64
+from io import BytesIO
+from models.image_text_transformation import ImageTextTransformation
+
+def pil_image_to_base64(image):
+ buffered = BytesIO()
+ image.save(buffered, format="JPEG")
+ img_str = base64.b64encode(buffered.getvalue()).decode()
+ return img_str
+
+def add_logo():
+ with open("examples/logo.png", "rb") as f:
+ logo_base64 = base64.b64encode(f.read()).decode()
+ return logo_base64
+
+def process_image(image_src, processor):
+ gen_text = processor.image_to_text(image_src)
+ gen_image = processor.text_to_image(gen_text)
+ gen_image_str = pil_image_to_base64(gen_image)
+ # Combine the outputs into a single HTML output
+ custom_output = f'''
+ {gen_text}
+
+ {gen_text}
+
+ {gen_text}
+
+ Understanding Image with Text'
+
+# Create Gradio interface
+interface = gr.Interface(
+ fn=lambda image: process_image(image, processor), # Pass the processor object using a lambda function
+ inputs=image_input,
+ outputs=gr.outputs.HTML(),
+ title=title_with_logo,
+ description="""
+ This code support image to text transformation. Then the generated text can do retrieval, question answering et al to conduct zero-shot.
+ """
+)
+
+# Launch the interface
+interface.launch()
\ No newline at end of file
diff --git a/models/__pycache__/blip2_model.cpython-38.pyc b/models/__pycache__/blip2_model.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..819e5a4cc1d62bd8a9865e946c75d6d65045c1be
Binary files /dev/null and b/models/__pycache__/blip2_model.cpython-38.pyc differ
diff --git a/models/__pycache__/blip2_model.cpython-39.pyc b/models/__pycache__/blip2_model.cpython-39.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..079f35e679cd1c43699ad374d70b89834565bf15
Binary files /dev/null and b/models/__pycache__/blip2_model.cpython-39.pyc differ
diff --git a/models/__pycache__/controlnet_model.cpython-38.pyc b/models/__pycache__/controlnet_model.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..dbaa59e91d6a7267e419a93426219cc56eae9ff5
Binary files /dev/null and b/models/__pycache__/controlnet_model.cpython-38.pyc differ
diff --git a/models/__pycache__/gpt_model.cpython-38.pyc b/models/__pycache__/gpt_model.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3a5644f28bfd1954d86e3523aee9f7fda8d0c975
Binary files /dev/null and b/models/__pycache__/gpt_model.cpython-38.pyc differ
diff --git a/models/__pycache__/grit_model.cpython-38.pyc b/models/__pycache__/grit_model.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6e79a2b8a67c228da3d02df6ec42c728a206d0db
Binary files /dev/null and b/models/__pycache__/grit_model.cpython-38.pyc differ
diff --git a/models/__pycache__/image_text_transformation.cpython-38.pyc b/models/__pycache__/image_text_transformation.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..acca6a68c1114aabe3042fc1aefc21375d64789d
Binary files /dev/null and b/models/__pycache__/image_text_transformation.cpython-38.pyc differ
diff --git a/models/__pycache__/image_text_transformation.cpython-39.pyc b/models/__pycache__/image_text_transformation.cpython-39.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6e8ff07c9bd604094179b83188bd2355813d59b9
Binary files /dev/null and b/models/__pycache__/image_text_transformation.cpython-39.pyc differ
diff --git a/models/__pycache__/region_semantic.cpython-38.pyc b/models/__pycache__/region_semantic.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..190cd00e72697ab1bfb1628b3919dcfd2f00e0c7
Binary files /dev/null and b/models/__pycache__/region_semantic.cpython-38.pyc differ
diff --git a/models/blip2_model.py b/models/blip2_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..afc6297f4d798f42ce8f172880111adc4a7490ac
--- /dev/null
+++ b/models/blip2_model.py
@@ -0,0 +1,38 @@
+from PIL import Image
+import requests
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
+import torch
+
+
+class ImageCaptioning:
+ def __init__(self) -> None:
+ self.device = None
+ # self.processor, self.model = None, None
+ self.processor, self.model = self.initialize_model()
+
+ def initialize_model(self):
+ # device = "cuda" if torch.cuda.is_available() else "cpu"
+ self.device = "cpu" # for low gpu memory devices
+ if self.device == 'cpu':
+ self.data_type = torch.float32
+ else:
+ self.data_type = torch.float16
+ processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+ model = Blip2ForConditionalGeneration.from_pretrained(
+ "Salesforce/blip2-opt-2.7b", torch_dtype=self.data_type
+ )
+ model.to(self.device)
+ return processor, model
+
+ def image_caption(self, image_src):
+ image = Image.open(image_src)
+ inputs = self.processor(images=image, return_tensors="pt").to(self.device, self.data_type)
+ generated_ids = self.model.generate(**inputs)
+ generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+ print('*'*100 + '\nStep1, BLIP2 caption:')
+ print(generated_text)
+ print('\n' + '*'*100)
+ return generated_text
+
+ def image_caption_debug(self, image_src):
+ return "A dish with salmon, broccoli, and something yellow."
\ No newline at end of file
diff --git a/models/controlnet_model.py b/models/controlnet_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..bac2bad34c77244350d5f28c620dc768734a34f2
--- /dev/null
+++ b/models/controlnet_model.py
@@ -0,0 +1,51 @@
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+from diffusers import (
+ StableDiffusionControlNetPipeline,
+ ControlNetModel,
+ UniPCMultistepScheduler,
+)
+
+
+class TextToImage:
+ def __init__(self):
+ # self.model = None
+ self.model = self.initialize_model()
+
+ def initialize_model(self):
+ controlnet = ControlNetModel.from_pretrained(
+ "fusing/stable-diffusion-v1-5-controlnet-canny",
+ torch_dtype=torch.float16,
+ )
+ pipeline = StableDiffusionControlNetPipeline.from_pretrained(
+ "runwayml/stable-diffusion-v1-5",
+ controlnet=controlnet,
+ safety_checker=None,
+ torch_dtype=torch.float16,
+ )
+ pipeline.scheduler = UniPCMultistepScheduler.from_config(
+ pipeline.scheduler.config
+ )
+ pipeline.enable_model_cpu_offload()
+ return pipeline
+
+ @staticmethod
+ def preprocess_image(image):
+ image = np.array(image)
+ low_threshold = 100
+ high_threshold = 200
+ image = cv2.Canny(image, low_threshold, high_threshold)
+ image = np.stack([image, image, image], axis=2)
+ image = Image.fromarray(image)
+ return image
+
+ def text_to_image(self, text, image):
+ image = self.preprocess_image(image)
+ generated_image = self.model(text, image, num_inference_steps=20).images[0]
+ return generated_image
+
+ def text_to_image_debug(self, text, image):
+ print("text_to_image_debug")
+ return image
\ No newline at end of file
diff --git a/models/gpt_model.py b/models/gpt_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..f4321497d52324048570c40a6e5c0b9bd3b99958
--- /dev/null
+++ b/models/gpt_model.py
@@ -0,0 +1,40 @@
+import openai
+
+class ImageToText:
+ def __init__(self, api_key):
+ self.template = self.initialize_template()
+ openai.api_key = api_key
+
+ def initialize_template(self):
+ prompt_prefix_1 = """Generate only an informative and nature paragraph based on the given information(a,b,c,d):\n"""
+ prompt_prefix_2 = """\n a. Image Resolution: """
+ prompt_prefix_3 = """\n b. Image Caption: """
+ prompt_prefix_4 = """\n c. Dense Caption: """
+ prompt_prefix_5 = """\n d. Region Semantic: """
+ prompt_suffix = """\n There are some rules:
+ Show object, color and position.
+ Use nouns rather than coordinates to show position information of each object.
+ No more than 7 sentences.
+ Only use one paragraph.
+ Do not appear number.
+ """
+ template = f"{prompt_prefix_1}{prompt_prefix_2}{{width}}X{{height}}{prompt_prefix_3}{{caption}}{prompt_prefix_4}{{dense_caption}}{prompt_prefix_5}{{region_semantic}}{prompt_suffix}"
+ return template
+
+ def paragraph_summary_with_gpt(self, caption, dense_caption, region_semantic, width, height):
+ question = self.template.format(width=width, height=height, caption=caption, dense_caption=dense_caption, region_semantic=region_semantic)
+ print('*'*100)
+ print("question:", question)
+ completion = openai.ChatCompletion.create(
+ model="gpt-3.5-turbo",
+ messages = [
+ {"role": "user", "content" : question}]
+ )
+ print("chatgpt response:", completion['choices'][0]['message']['content'])
+ print('*'*100)
+ return completion['choices'][0]['message']['content']
+
+ def paragraph_summary_with_gpt_debug(self, caption, dense_caption, width, height):
+ question = self.template.format(width=width, height=height, caption=caption, dense_caption=dense_caption)
+ print("paragraph_summary_with_gpt_debug:")
+ return question
diff --git a/models/grit_model.py b/models/grit_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..3634eebaf0dcb4c03801649641aedb42db285ca4
--- /dev/null
+++ b/models/grit_model.py
@@ -0,0 +1,26 @@
+import os
+from models.grit_src.image_dense_captions import image_caption_api
+
+class DenseCaptioning():
+ def __init__(self) -> None:
+ self.model = None
+
+
+ def initialize_model(self):
+ pass
+
+ def image_dense_caption_debug(self, image_src):
+ dense_caption = """
+ 1. the broccoli is green, [0, 0, 333, 325];
+ 2. a piece of broccoli, [0, 147, 143, 324];
+ 3. silver fork on plate, [4, 547, 252, 612];
+ """
+ return dense_caption
+
+ def image_dense_caption(self, image_src):
+ dense_caption = image_caption_api(image_src)
+ print("Step2, Dense Caption:\n")
+ print(dense_caption)
+ print('\n'+'*'*100)
+ return dense_caption
+
\ No newline at end of file
diff --git a/models/grit_src/__pycache__/image_dense_captions.cpython-38.pyc b/models/grit_src/__pycache__/image_dense_captions.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e47d2ae24c8d5b10f44f411e450f0e28d5e5bc3f
Binary files /dev/null and b/models/grit_src/__pycache__/image_dense_captions.cpython-38.pyc differ
diff --git a/models/grit_src/configs/Base.yaml b/models/grit_src/configs/Base.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..445690acaafacfba6b54f28b4cf32e40c4bcae9d
--- /dev/null
+++ b/models/grit_src/configs/Base.yaml
@@ -0,0 +1,77 @@
+MODEL:
+ META_ARCHITECTURE: "GRiT"
+ MASK_ON: True
+ PROPOSAL_GENERATOR:
+ NAME: "CenterNet"
+ FPN:
+ IN_FEATURES: ["layer3", "layer4", "layer5"]
+ PIXEL_MEAN: [123.675, 116.280, 103.530]
+ PIXEL_STD: [58.395, 57.12, 57.375]
+ ROI_HEADS:
+ NAME: GRiTROIHeadsAndTextDecoder
+ IN_FEATURES: ["p3", "p4", "p5"]
+ IOU_THRESHOLDS: [0.6]
+ NUM_CLASSES: 1
+ SCORE_THRESH_TEST: 0.02
+ NMS_THRESH_TEST: 0.5
+ OBJECT_FEAT_POOLER_RES: 14
+ ROI_BOX_CASCADE_HEAD:
+ IOUS: [0.6, 0.7, 0.8]
+ ROI_BOX_HEAD:
+ NAME: "FastRCNNConvFCHead"
+ NUM_FC: 2
+ POOLER_RESOLUTION: 7
+ CLS_AGNOSTIC_BBOX_REG: True
+ MULT_PROPOSAL_SCORE: True
+ ROI_MASK_HEAD:
+ NAME: "MaskRCNNConvUpsampleHead"
+ NUM_CONV: 4
+ POOLER_RESOLUTION: 14
+ CLS_AGNOSTIC_MASK: True
+ CENTERNET:
+ NUM_CLASSES: 1
+ REG_WEIGHT: 1.
+ NOT_NORM_REG: True
+ ONLY_PROPOSAL: True
+ WITH_AGN_HM: True
+ INFERENCE_TH: 0.0001
+ PRE_NMS_TOPK_TRAIN: 4000
+ POST_NMS_TOPK_TRAIN: 2000
+ PRE_NMS_TOPK_TEST: 1000
+ POST_NMS_TOPK_TEST: 256
+ NMS_TH_TRAIN: 0.9
+ NMS_TH_TEST: 0.9
+ POS_WEIGHT: 0.5
+ NEG_WEIGHT: 0.5
+ IGNORE_HIGH_FP: 0.85
+DATASETS:
+ TRAIN: ("coco_2017_train",)
+ TEST: ("coco_2017_val",)
+DATALOADER:
+ SAMPLER_TRAIN: "MultiDatasetSampler"
+ DATASET_RATIO: [1]
+ DATASET_INPUT_SIZE: [1024]
+ DATASET_INPUT_SCALE: [[0.1, 2.0]]
+ FILTER_EMPTY_ANNOTATIONS: False
+ NUM_WORKERS: 8
+TEST:
+ DETECTIONS_PER_IMAGE: 256
+SOLVER:
+ LR_SCHEDULER_NAME: "WarmupCosineLR"
+ CHECKPOINT_PERIOD: 10000
+ WARMUP_ITERS: 1000
+ WARMUP_FACTOR: 0.001
+ USE_CUSTOM_SOLVER: True
+ OPTIMIZER: "ADAMW"
+ MAX_ITER: 180000
+ IMS_PER_BATCH: 64
+ BASE_LR: 0.00008
+ VIT_LAYER_DECAY: True
+ CLIP_GRADIENTS:
+ ENABLED: True
+INPUT:
+ FORMAT: RGB
+ CUSTOM_AUG: EfficientDetResizeCrop
+ TRAIN_SIZE: 640
+USE_ACT_CHECKPOINT: True
+VERSION: 2
\ No newline at end of file
diff --git a/models/grit_src/configs/GRiT_B_DenseCap.yaml b/models/grit_src/configs/GRiT_B_DenseCap.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..0e7d2d2c7448d330d9356b3af90b975b2ce7d528
--- /dev/null
+++ b/models/grit_src/configs/GRiT_B_DenseCap.yaml
@@ -0,0 +1,20 @@
+_BASE_: "Base.yaml"
+MODEL:
+ TRAIN_TASK: ["DenseCap"]
+ TEST_TASK: "DenseCap"
+ MASK_ON: False
+ ROI_HEADS:
+ SOFT_NMS_ENABLED: False
+ BEAM_SIZE: 1
+ WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth"
+ BACKBONE:
+ NAME: build_vit_fpn_backbone
+ VIT_LAYERS: 12
+SOLVER:
+ VIT_LAYER_DECAY_RATE: 0.7
+DATASETS:
+ TRAIN: ("vg_train",)
+ TEST: ("vg_test",)
+DATALOADER:
+ DATASET_BS: 2
+OUTPUT_DIR: "./output/GRiT_B_DenseCap"
\ No newline at end of file
diff --git a/models/grit_src/configs/GRiT_B_DenseCap_ObjectDet.yaml b/models/grit_src/configs/GRiT_B_DenseCap_ObjectDet.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..49f3ef13ab8bf0eb8515c009e70e1d33687efd39
--- /dev/null
+++ b/models/grit_src/configs/GRiT_B_DenseCap_ObjectDet.yaml
@@ -0,0 +1,23 @@
+_BASE_: "Base.yaml"
+MODEL:
+ TRAIN_TASK: ["ObjectDet", "DenseCap"]
+ TEST_TASK: "DenseCap" # DenseCap or ObjectDet: Choose one for testing
+ MASK_ON: True
+ ROI_HEADS:
+ SOFT_NMS_ENABLED: False
+ BEAM_SIZE: 1
+ WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth"
+ BACKBONE:
+ NAME: build_vit_fpn_backbone
+ VIT_LAYERS: 12
+SOLVER:
+ VIT_LAYER_DECAY_RATE: 0.7
+DATASETS:
+ TRAIN: ("GRiT_coco2017_train", "vg_train")
+ TEST: ("coco_2017_test-dev",)
+DATALOADER:
+ DATASET_RATIO: [1, 1]
+ DATASET_BS: 2
+ DATASET_INPUT_SIZE: [1024, 1024]
+ DATASET_INPUT_SCALE: [[0.1, 2.0], [0.1, 2.0]]
+OUTPUT_DIR: "./output/GRiT_B_DenseCap_ObjectDet"
\ No newline at end of file
diff --git a/models/grit_src/configs/GRiT_B_ObjectDet.yaml b/models/grit_src/configs/GRiT_B_ObjectDet.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..e7a75052f84b7913480cc5ca0e29c03e4dbea4ef
--- /dev/null
+++ b/models/grit_src/configs/GRiT_B_ObjectDet.yaml
@@ -0,0 +1,20 @@
+_BASE_: "Base.yaml"
+MODEL:
+ TRAIN_TASK: ["ObjectDet"]
+ TEST_TASK: "ObjectDet"
+ MASK_ON: True
+ ROI_HEADS:
+ SOFT_NMS_ENABLED: True
+ BEAM_SIZE: 3
+ WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth"
+ BACKBONE:
+ NAME: build_vit_fpn_backbone
+ VIT_LAYERS: 12
+SOLVER:
+ VIT_LAYER_DECAY_RATE: 0.7
+DATASETS:
+ TRAIN: ("GRiT_coco2017_train",)
+ TEST: ("coco_2017_val",)
+DATALOADER:
+ DATASET_BS: 2
+OUTPUT_DIR: "./output/GRiT_B_ObjectDet"
\ No newline at end of file
diff --git a/models/grit_src/configs/GRiT_H_ObjectDet.yaml b/models/grit_src/configs/GRiT_H_ObjectDet.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..000a1d4629b695f5c4b4741fe28d0b8561c11cdb
--- /dev/null
+++ b/models/grit_src/configs/GRiT_H_ObjectDet.yaml
@@ -0,0 +1,21 @@
+_BASE_: "Base.yaml"
+MODEL:
+ TRAIN_TASK: ["ObjectDet"]
+ TEST_TASK: "ObjectDet"
+ MASK_ON: True
+ ROI_HEADS:
+ SOFT_NMS_ENABLED: True
+ BEAM_SIZE: 3
+ WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_huge_p14to16.pth"
+ BACKBONE:
+ NAME: build_vit_fpn_backbone_huge
+ VIT_LAYERS: 32
+SOLVER:
+ MAX_ITER: 135000
+ VIT_LAYER_DECAY_RATE: 0.9
+DATASETS:
+ TRAIN: ("GRiT_coco2017_train",)
+ TEST: ("coco_2017_val",)
+DATALOADER:
+ DATASET_BS: 1
+OUTPUT_DIR: "./output/GRiT_H_ObjectDet"
\ No newline at end of file
diff --git a/models/grit_src/configs/GRiT_L_ObjectDet.yaml b/models/grit_src/configs/GRiT_L_ObjectDet.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..b6e3b97f08fe4671e1a686b6cb6a83f8fc52f9a7
--- /dev/null
+++ b/models/grit_src/configs/GRiT_L_ObjectDet.yaml
@@ -0,0 +1,20 @@
+_BASE_: "Base.yaml"
+MODEL:
+ TRAIN_TASK: ["ObjectDet"]
+ TEST_TASK: "ObjectDet"
+ MASK_ON: True
+ ROI_HEADS:
+ SOFT_NMS_ENABLED: True
+ BEAM_SIZE: 3
+ WEIGHTS: "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_large.pth"
+ BACKBONE:
+ NAME: build_vit_fpn_backbone_large
+ VIT_LAYERS: 24
+SOLVER:
+ VIT_LAYER_DECAY_RATE: 0.8
+DATASETS:
+ TRAIN: ("GRiT_coco2017_train",)
+ TEST: ("coco_2017_val",)
+DATALOADER:
+ DATASET_BS: 1
+OUTPUT_DIR: "./output/GRiT_L_ObjectDet"
\ No newline at end of file
diff --git a/models/grit_src/grit/__init__.py b/models/grit_src/grit/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..81f24566b0093edc133440090715b20ee569ca37
--- /dev/null
+++ b/models/grit_src/grit/__init__.py
@@ -0,0 +1,7 @@
+from .modeling.meta_arch import grit
+from .modeling.roi_heads import grit_roi_heads
+from .modeling.backbone import vit
+
+from .data.datasets import object365
+from .data.datasets import vg
+from .data.datasets import grit_coco
\ No newline at end of file
diff --git a/models/grit_src/grit/__pycache__/__init__.cpython-38.pyc b/models/grit_src/grit/__pycache__/__init__.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..132bc0316cd3edc52ce900691fc877e72e41b3e1
Binary files /dev/null and b/models/grit_src/grit/__pycache__/__init__.cpython-38.pyc differ
diff --git a/models/grit_src/grit/__pycache__/config.cpython-38.pyc b/models/grit_src/grit/__pycache__/config.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ca36ee0b9ebd33477803efb0dc31fc62515bb604
Binary files /dev/null and b/models/grit_src/grit/__pycache__/config.cpython-38.pyc differ
diff --git a/models/grit_src/grit/__pycache__/predictor.cpython-38.pyc b/models/grit_src/grit/__pycache__/predictor.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1e7119650058605d0929c5b71622565a7e7e9c9f
Binary files /dev/null and b/models/grit_src/grit/__pycache__/predictor.cpython-38.pyc differ
diff --git a/models/grit_src/grit/config.py b/models/grit_src/grit/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..fabe7f0fbe1e41c6eb280f8f7d6ae2e9c4911135
--- /dev/null
+++ b/models/grit_src/grit/config.py
@@ -0,0 +1,50 @@
+from detectron2.config import CfgNode as CN
+
+
+def add_grit_config(cfg):
+ _C = cfg
+
+ _C.MODEL.BEAM_SIZE = 1
+ _C.MODEL.TRAIN_TASK = ["ObjectDet", "DenseCap"]
+ _C.MODEL.TEST_TASK = "DenseCap" # This can be varied if the model is jointly trained on multiple tasks
+
+ _C.MODEL.ROI_BOX_HEAD.USE_BIAS = 0.0 # >= 0: not use
+ _C.MODEL.ROI_BOX_HEAD.MULT_PROPOSAL_SCORE = False
+
+ _C.MODEL.ROI_HEADS.MASK_WEIGHT = 1.0
+ _C.MODEL.ROI_HEADS.OBJECT_FEAT_POOLER_RES = 14
+ _C.MODEL.ROI_HEADS.SOFT_NMS_ENABLED = False
+
+ # Backbones
+ _C.MODEL.VIT_LAYERS = 12
+
+ # Text Decoder
+ _C.TEXT_DECODER = CN()
+ _C.TEXT_DECODER.VOCAB_SIZE = 30522
+ _C.TEXT_DECODER.HIDDEN_SIZE = 768
+ _C.TEXT_DECODER.NUM_LAYERS = 6
+ _C.TEXT_DECODER.ATTENTION_HEADS = 12
+ _C.TEXT_DECODER.FEEDFORWARD_SIZE = 768 * 4
+
+ # Multi-dataset dataloader
+ _C.DATALOADER.DATASET_RATIO = [1, 1] # sample ratio
+ _C.DATALOADER.DATASET_BS = 1
+ _C.DATALOADER.DATASET_INPUT_SIZE = [1024, 1024]
+ _C.DATALOADER.DATASET_INPUT_SCALE = [(0.1, 2.0), (0.1, 2.0)]
+ _C.DATALOADER.DATASET_MIN_SIZES = [(640, 800), (640, 800)]
+ _C.DATALOADER.DATASET_MAX_SIZES = [1333, 1333]
+
+ _C.SOLVER.USE_CUSTOM_SOLVER = True
+ _C.SOLVER.OPTIMIZER = 'ADAMW'
+ _C.SOLVER.VIT_LAYER_DECAY = True
+ _C.SOLVER.VIT_LAYER_DECAY_RATE = 0.7
+
+ _C.INPUT.CUSTOM_AUG = 'EfficientDetResizeCrop'
+ _C.INPUT.TRAIN_SIZE = 1024
+ _C.INPUT.TEST_SIZE = 1024
+ _C.INPUT.SCALE_RANGE = (0.1, 2.)
+ # 'default' for fixed short / long edge
+ _C.INPUT.TEST_INPUT_TYPE = 'default'
+
+ _C.FIND_UNUSED_PARAM = True
+ _C.USE_ACT_CHECKPOINT = True
\ No newline at end of file
diff --git a/models/grit_src/grit/custom_solver.py b/models/grit_src/grit/custom_solver.py
new file mode 100644
index 0000000000000000000000000000000000000000..87f7d61ed756acf9326b7ab4097a989a9e6c7532
--- /dev/null
+++ b/models/grit_src/grit/custom_solver.py
@@ -0,0 +1,88 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Modified by Jialian Wu from https://github.com/facebookresearch/Detic/blob/main/detic/custom_solver.py
+import itertools
+from typing import Any, Callable, Dict, Iterable, List, Set, Type, Union
+import torch
+
+from detectron2.config import CfgNode
+
+from detectron2.solver.build import maybe_add_gradient_clipping
+
+
+def build_custom_optimizer(cfg: CfgNode, model: torch.nn.Module) -> torch.optim.Optimizer:
+ params: List[Dict[str, Any]] = []
+ memo: Set[torch.nn.parameter.Parameter] = set()
+ optimizer_type = cfg.SOLVER.OPTIMIZER
+
+ for key, value in model.named_parameters(recurse=True):
+ if not value.requires_grad:
+ continue
+ # Avoid duplicating parameters
+ if value in memo:
+ continue
+ memo.add(value)
+ lr = cfg.SOLVER.BASE_LR
+ weight_decay = cfg.SOLVER.WEIGHT_DECAY
+
+ if cfg.SOLVER.VIT_LAYER_DECAY:
+ lr = lr * get_vit_lr_decay_rate(key, cfg.SOLVER.VIT_LAYER_DECAY_RATE, cfg.MODEL.VIT_LAYERS)
+
+ param = {"params": [value], "lr": lr}
+ if optimizer_type != 'ADAMW':
+ param['weight_decay'] = weight_decay
+ params += [param]
+
+ def maybe_add_full_model_gradient_clipping(optim): # optim: the optimizer class
+ # detectron2 doesn't have full model gradient clipping now
+ clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE
+ enable = (
+ cfg.SOLVER.CLIP_GRADIENTS.ENABLED
+ and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model"
+ and clip_norm_val > 0.0
+ )
+
+ class FullModelGradientClippingOptimizer(optim):
+ def step(self, closure=None):
+ all_params = itertools.chain(*[x["params"] for x in self.param_groups])
+ torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
+ super().step(closure=closure)
+
+ return FullModelGradientClippingOptimizer if enable else optim
+
+
+ if optimizer_type == 'SGD':
+ optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(
+ params, cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM,
+ nesterov=cfg.SOLVER.NESTEROV
+ )
+ elif optimizer_type == 'ADAMW':
+ optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(
+ params, cfg.SOLVER.BASE_LR,
+ weight_decay=cfg.SOLVER.WEIGHT_DECAY
+ )
+ else:
+ raise NotImplementedError(f"no optimizer type {optimizer_type}")
+ if not cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model":
+ optimizer = maybe_add_gradient_clipping(cfg, optimizer)
+ return optimizer
+
+
+def get_vit_lr_decay_rate(name, lr_decay_rate=1.0, num_layers=12):
+ """
+ Calculate lr decay rate for different ViT blocks.
+ Args:
+ name (string): parameter name.
+ lr_decay_rate (float): base lr decay rate.
+ num_layers (int): number of ViT blocks.
+
+ Returns:
+ lr decay rate for the given parameter.
+ """
+ layer_id = num_layers + 1
+ if name.startswith("backbone"):
+ if ".pos_embed" in name or ".patch_embed" in name:
+ layer_id = 0
+ elif ".blocks." in name and ".residual." not in name:
+ layer_id = int(name[name.find(".blocks.") :].split(".")[2]) + 1
+
+ return lr_decay_rate ** (num_layers + 1 - layer_id)
\ No newline at end of file
diff --git a/models/grit_src/grit/data/__pycache__/custom_build_augmentation.cpython-38.pyc b/models/grit_src/grit/data/__pycache__/custom_build_augmentation.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fb14a971188efed0f1ebc1023edfed3a1afa26f7
Binary files /dev/null and b/models/grit_src/grit/data/__pycache__/custom_build_augmentation.cpython-38.pyc differ
diff --git a/models/grit_src/grit/data/__pycache__/custom_dataset_mapper.cpython-38.pyc b/models/grit_src/grit/data/__pycache__/custom_dataset_mapper.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cce6aabfaddc99e97e0d6b542fb31d7dcb884328
Binary files /dev/null and b/models/grit_src/grit/data/__pycache__/custom_dataset_mapper.cpython-38.pyc differ
diff --git a/models/grit_src/grit/data/custom_build_augmentation.py b/models/grit_src/grit/data/custom_build_augmentation.py
new file mode 100644
index 0000000000000000000000000000000000000000..49a52d011c09dbe027d41ee7e50127c392a8bf33
--- /dev/null
+++ b/models/grit_src/grit/data/custom_build_augmentation.py
@@ -0,0 +1,44 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from detectron2.data import transforms as T
+from .transforms.custom_augmentation_impl import EfficientDetResizeCrop
+
+
+def build_custom_augmentation(cfg, is_train, scale=None, size=None, \
+ min_size=None, max_size=None):
+ """
+ Create a list of default :class:`Augmentation` from config.
+ Now it includes resizing and flipping.
+
+ Returns:
+ list[Augmentation]
+ """
+ if cfg.INPUT.CUSTOM_AUG == 'ResizeShortestEdge':
+ if is_train:
+ min_size = cfg.INPUT.MIN_SIZE_TRAIN if min_size is None else min_size
+ max_size = cfg.INPUT.MAX_SIZE_TRAIN if max_size is None else max_size
+ sample_style = cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING
+ else:
+ min_size = cfg.INPUT.MIN_SIZE_TEST
+ max_size = cfg.INPUT.MAX_SIZE_TEST
+ sample_style = "choice"
+ augmentation = [T.ResizeShortestEdge(min_size, max_size, sample_style)]
+ elif cfg.INPUT.CUSTOM_AUG == 'EfficientDetResizeCrop':
+ if is_train:
+ scale = cfg.INPUT.SCALE_RANGE if scale is None else scale
+ size = cfg.INPUT.TRAIN_SIZE if size is None else size
+ else:
+ scale = (1, 1)
+ size = cfg.INPUT.TEST_SIZE
+ augmentation = [EfficientDetResizeCrop(size, scale)]
+ else:
+ assert 0, cfg.INPUT.CUSTOM_AUG
+
+ if is_train:
+ augmentation.append(T.RandomFlip())
+ return augmentation
+
+
+build_custom_transform_gen = build_custom_augmentation
+"""
+Alias for backward-compatibility.
+"""
\ No newline at end of file
diff --git a/models/grit_src/grit/data/custom_dataset_dataloader.py b/models/grit_src/grit/data/custom_dataset_dataloader.py
new file mode 100644
index 0000000000000000000000000000000000000000..ea9c4172f838d130df297bed9c0755669720c39d
--- /dev/null
+++ b/models/grit_src/grit/data/custom_dataset_dataloader.py
@@ -0,0 +1,250 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Jialian Wu from https://github.com/facebookresearch/Detic/blob/main/detic/data/custom_dataset_dataloader.py
+import operator
+import torch
+import torch.utils.data
+from detectron2.utils.comm import get_world_size
+
+from detectron2.config import configurable
+from torch.utils.data.sampler import BatchSampler, Sampler
+from detectron2.data.common import DatasetFromList, MapDataset
+from detectron2.data.dataset_mapper import DatasetMapper
+from detectron2.data.build import get_detection_dataset_dicts, build_batch_data_loader
+from detectron2.data.samplers import TrainingSampler
+from detectron2.data.build import worker_init_reset_seed, print_instances_class_histogram
+from detectron2.data.build import filter_images_with_only_crowd_annotations
+from detectron2.data.build import filter_images_with_few_keypoints
+from detectron2.data.build import check_metadata_consistency
+from detectron2.data.catalog import MetadataCatalog, DatasetCatalog
+from detectron2.utils import comm
+import itertools
+from typing import Optional
+
+
+def _custom_train_loader_from_config(cfg, mapper=None, *, dataset=None, sampler=None):
+ sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
+ if 'MultiDataset' in sampler_name:
+ dataset_dicts = get_detection_dataset_dicts_with_source(
+ cfg.DATASETS.TRAIN,
+ filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
+ min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
+ if cfg.MODEL.KEYPOINT_ON else 0,
+ proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
+ )
+ else:
+ dataset_dicts = get_detection_dataset_dicts(
+ cfg.DATASETS.TRAIN,
+ filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
+ min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
+ if cfg.MODEL.KEYPOINT_ON else 0,
+ proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
+ )
+
+ if mapper is None:
+ mapper = DatasetMapper(cfg, True)
+
+ if sampler is not None:
+ pass
+ elif sampler_name == "TrainingSampler":
+ sampler = TrainingSampler(len(dataset))
+ elif sampler_name == "MultiDatasetSampler":
+ sampler = MultiDatasetSampler(
+ dataset_dicts,
+ dataset_ratio=cfg.DATALOADER.DATASET_RATIO,
+ )
+ else:
+ raise ValueError("Unknown training sampler: {}".format(sampler_name))
+
+ return {
+ "dataset": dataset_dicts,
+ "sampler": sampler,
+ "mapper": mapper,
+ "total_batch_size": cfg.SOLVER.IMS_PER_BATCH,
+ "num_workers": cfg.DATALOADER.NUM_WORKERS,
+ 'dataset_bs': cfg.DATALOADER.DATASET_BS,
+ 'num_datasets': len(cfg.DATASETS.TRAIN)
+ }
+
+
+@configurable(from_config=_custom_train_loader_from_config)
+def build_custom_train_loader(
+ dataset, *, mapper, sampler,
+ total_batch_size=16,
+ num_workers=0,
+ num_datasets=1,
+ dataset_bs=1
+):
+
+ if isinstance(dataset, list):
+ dataset = DatasetFromList(dataset, copy=False)
+ if mapper is not None:
+ dataset = MapDataset(dataset, mapper)
+ if sampler is None:
+ sampler = TrainingSampler(len(dataset))
+ assert isinstance(sampler, torch.utils.data.sampler.Sampler)
+
+ return build_dataset_batch_data_loader(
+ dataset_bs,
+ dataset,
+ sampler,
+ total_batch_size,
+ num_datasets=num_datasets,
+ num_workers=num_workers,
+ )
+
+
+def build_dataset_batch_data_loader(
+ dataset_bs, dataset, sampler, total_batch_size, num_datasets, num_workers=0
+):
+
+ world_size = get_world_size()
+ assert (
+ total_batch_size > 0 and total_batch_size % world_size == 0
+ ), "Total batch size ({}) must be divisible by the number of gpus ({}).".format(
+ total_batch_size, world_size
+ )
+
+ data_loader = torch.utils.data.DataLoader(
+ dataset,
+ sampler=sampler,
+ num_workers=num_workers,
+ batch_sampler=None,
+ collate_fn=operator.itemgetter(0), # don't batch, but yield individual elements
+ worker_init_fn=worker_init_reset_seed,
+ )
+
+ if num_datasets > 1:
+ return MultiDatasets(data_loader, dataset_bs, num_datasets)
+ else:
+ return SingleDataset(data_loader, dataset_bs)
+
+
+def get_detection_dataset_dicts_with_source(
+ dataset_names, filter_empty=True, min_keypoints=0, proposal_files=None
+):
+ assert len(dataset_names)
+ dataset_dicts = [DatasetCatalog.get(dataset_name) for dataset_name in dataset_names]
+ for dataset_name, dicts in zip(dataset_names, dataset_dicts):
+ assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
+
+ for source_id, (dataset_name, dicts) in \
+ enumerate(zip(dataset_names, dataset_dicts)):
+ assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
+ for d in dicts:
+ d['dataset_source'] = source_id
+
+ if "annotations" in dicts[0]:
+ try:
+ class_names = MetadataCatalog.get(dataset_name).thing_classes
+ check_metadata_consistency("thing_classes", dataset_name)
+ print_instances_class_histogram(dicts, class_names)
+ except AttributeError: # class names are not available for this dataset
+ pass
+
+ assert proposal_files is None
+
+ dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))
+
+ has_instances = "annotations" in dataset_dicts[0]
+ if filter_empty and has_instances:
+ dataset_dicts = filter_images_with_only_crowd_annotations(dataset_dicts)
+ if min_keypoints > 0 and has_instances:
+ dataset_dicts = filter_images_with_few_keypoints(dataset_dicts, min_keypoints)
+
+ return dataset_dicts
+
+
+class MultiDatasetSampler(Sampler):
+ def __init__(
+ self,
+ dataset_dicts,
+ dataset_ratio,
+ seed: Optional[int] = None,
+ ):
+ sizes = [0 for _ in range(len(dataset_ratio))]
+ for d in dataset_dicts:
+ sizes[d['dataset_source']] += 1
+ print('dataset sizes', sizes)
+ self.sizes = sizes
+ assert len(dataset_ratio) == len(sizes), \
+ 'length of dataset ratio {} should be equal to number if dataset {}'.format(
+ len(dataset_ratio), len(sizes)
+ )
+ if seed is None:
+ seed = comm.shared_random_seed()
+ self._seed = int(seed)
+ self._rank = comm.get_rank()
+ self._world_size = comm.get_world_size()
+
+ self.dataset_ids = torch.tensor(
+ [d['dataset_source'] for d in dataset_dicts], dtype=torch.long)
+ self.dataset_ratio = dataset_ratio
+
+ dataset_weight = [torch.ones(s) * max(sizes) / s * r / sum(dataset_ratio) \
+ for i, (r, s) in enumerate(zip(dataset_ratio, sizes))]
+ dataset_weight = torch.cat(dataset_weight)
+
+ self.weights = dataset_weight
+ self.sample_epoch_size = len(self.weights)
+
+ def __iter__(self):
+ start = self._rank
+ yield from itertools.islice(
+ self._infinite_indices(), start, None, self._world_size)
+
+ def _infinite_indices(self):
+ g = torch.Generator()
+ g.manual_seed(self._seed)
+ while True:
+ if len(self.dataset_ratio) > 1:
+ # multiple datasets
+ ids = torch.multinomial(
+ self.weights, self.sample_epoch_size, generator=g,
+ replacement=True)
+ nums = [(self.dataset_ids[ids] == i).sum().int().item() \
+ for i in range(len(self.sizes))]
+ yield from ids
+ else:
+ # single dataset
+ yield from torch.randperm(self.sizes[0], generator=g).tolist()
+
+
+class SingleDataset(torch.utils.data.IterableDataset):
+ def __init__(self, dataset, batch_sizes):
+ self.dataset = dataset
+ self.batch_sizes = batch_sizes
+ self._buckets = [[] for _ in range(2)]
+
+ def __iter__(self):
+ for d in self.dataset:
+ w, h = d["width"], d["height"]
+ aspect_ratio_bucket_id = 0 if w > h else 1
+ bucket_id = aspect_ratio_bucket_id
+ bucket = self._buckets[bucket_id]
+ bucket.append(d)
+ if len(bucket) == self.batch_sizes:
+ yield bucket[:]
+ del bucket[:]
+
+
+class MultiDatasets(torch.utils.data.IterableDataset):
+ def __init__(self, dataset, batch_sizes, num_datasets):
+ self.dataset = dataset
+ self.batch_sizes = batch_sizes
+ self._buckets = [[] for _ in range(2 * num_datasets)]
+ self.iter_idx = 0
+ self.num_datasets = num_datasets
+
+ def __iter__(self):
+ for d in self.dataset:
+ w, h = d["width"], d["height"]
+ aspect_ratio_bucket_id = 0 if w > h else 1
+ bucket_id = d['dataset_source'] * 2 + aspect_ratio_bucket_id
+ bucket = self._buckets[bucket_id]
+ if len(bucket) < self.batch_sizes:
+ bucket.append(d)
+ selected_dataset = self.iter_idx % self.num_datasets
+ if len(bucket) == self.batch_sizes and selected_dataset == d['dataset_source']:
+ self.iter_idx += 1
+ yield bucket[:]
+ del bucket[:]
\ No newline at end of file
diff --git a/models/grit_src/grit/data/custom_dataset_mapper.py b/models/grit_src/grit/data/custom_dataset_mapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e21edb3d151dafdca5c4debfb7341a9ed0efdd9
--- /dev/null
+++ b/models/grit_src/grit/data/custom_dataset_mapper.py
@@ -0,0 +1,149 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Modified by Jialian Wu from https://github.com/facebookresearch/Detic/blob/main/detic/data/custom_dataset_mapper.py
+import copy
+import numpy as np
+import torch
+
+from detectron2.config import configurable
+
+from detectron2.data import detection_utils as utils
+from detectron2.data import transforms as T
+from detectron2.data.dataset_mapper import DatasetMapper
+from .custom_build_augmentation import build_custom_augmentation
+from itertools import compress
+import logging
+
+__all__ = ["CustomDatasetMapper", "ObjDescription"]
+logger = logging.getLogger(__name__)
+
+
+class CustomDatasetMapper(DatasetMapper):
+ @configurable
+ def __init__(self, is_train: bool,
+ dataset_augs=[],
+ **kwargs):
+ if is_train:
+ self.dataset_augs = [T.AugmentationList(x) for x in dataset_augs]
+ super().__init__(is_train, **kwargs)
+
+ @classmethod
+ def from_config(cls, cfg, is_train: bool = True):
+ ret = super().from_config(cfg, is_train)
+ if is_train:
+ if cfg.INPUT.CUSTOM_AUG == 'EfficientDetResizeCrop':
+ dataset_scales = cfg.DATALOADER.DATASET_INPUT_SCALE
+ dataset_sizes = cfg.DATALOADER.DATASET_INPUT_SIZE
+ ret['dataset_augs'] = [
+ build_custom_augmentation(cfg, True, scale, size) \
+ for scale, size in zip(dataset_scales, dataset_sizes)]
+ else:
+ assert cfg.INPUT.CUSTOM_AUG == 'ResizeShortestEdge'
+ min_sizes = cfg.DATALOADER.DATASET_MIN_SIZES
+ max_sizes = cfg.DATALOADER.DATASET_MAX_SIZES
+ ret['dataset_augs'] = [
+ build_custom_augmentation(
+ cfg, True, min_size=mi, max_size=ma) \
+ for mi, ma in zip(min_sizes, max_sizes)]
+ else:
+ ret['dataset_augs'] = []
+
+ return ret
+
+ def __call__(self, dataset_dict):
+ dataset_dict_out = self.prepare_data(dataset_dict)
+
+ # When augmented image is too small, do re-augmentation
+ retry = 0
+ while (dataset_dict_out["image"].shape[1] < 32 or dataset_dict_out["image"].shape[2] < 32):
+ retry += 1
+ if retry == 100:
+ logger.info('Retry 100 times for augmentation. Make sure the image size is not too small.')
+ logger.info('Find image information below')
+ logger.info(dataset_dict)
+ dataset_dict_out = self.prepare_data(dataset_dict)
+
+ return dataset_dict_out
+
+ def prepare_data(self, dataset_dict_in):
+ dataset_dict = copy.deepcopy(dataset_dict_in)
+ if 'file_name' in dataset_dict:
+ ori_image = utils.read_image(
+ dataset_dict["file_name"], format=self.image_format)
+ else:
+ ori_image, _, _ = self.tar_dataset[dataset_dict["tar_index"]]
+ ori_image = utils._apply_exif_orientation(ori_image)
+ ori_image = utils.convert_PIL_to_numpy(ori_image, self.image_format)
+ utils.check_image_size(dataset_dict, ori_image)
+
+ aug_input = T.AugInput(copy.deepcopy(ori_image), sem_seg=None)
+ if self.is_train:
+ transforms = \
+ self.dataset_augs[dataset_dict['dataset_source']](aug_input)
+ else:
+ transforms = self.augmentations(aug_input)
+ image, sem_seg_gt = aug_input.image, aug_input.sem_seg
+
+ image_shape = image.shape[:2]
+ dataset_dict["image"] = torch.as_tensor(
+ np.ascontiguousarray(image.transpose(2, 0, 1)))
+
+ if not self.is_train:
+ # USER: Modify this if you want to keep them for some reason.
+ dataset_dict.pop("annotations", None)
+ return dataset_dict
+
+ if "annotations" in dataset_dict:
+ if len(dataset_dict["annotations"]) > 0:
+ object_descriptions = [an['object_description'] for an in dataset_dict["annotations"]]
+ else:
+ object_descriptions = []
+ # USER: Modify this if you want to keep them for some reason.
+ for anno in dataset_dict["annotations"]:
+ if not self.use_instance_mask:
+ anno.pop("segmentation", None)
+ if not self.use_keypoint:
+ anno.pop("keypoints", None)
+
+ all_annos = [
+ (utils.transform_instance_annotations(
+ obj, transforms, image_shape,
+ keypoint_hflip_indices=self.keypoint_hflip_indices,
+ ), obj.get("iscrowd", 0))
+ for obj in dataset_dict.pop("annotations")
+ ]
+ annos = [ann[0] for ann in all_annos if ann[1] == 0]
+ instances = utils.annotations_to_instances(
+ annos, image_shape, mask_format=self.instance_mask_format
+ )
+
+ instances.gt_object_descriptions = ObjDescription(object_descriptions)
+
+ del all_annos
+ if self.recompute_boxes:
+ instances.gt_boxes = instances.gt_masks.get_bounding_boxes()
+ dataset_dict["instances"] = utils.filter_empty_instances(instances)
+
+ return dataset_dict
+
+
+class ObjDescription:
+ def __init__(self, object_descriptions):
+ self.data = object_descriptions
+
+ def __getitem__(self, item):
+ assert type(item) == torch.Tensor
+ assert item.dim() == 1
+ if len(item) > 0:
+ assert item.dtype == torch.int64 or item.dtype == torch.bool
+ if item.dtype == torch.int64:
+ return ObjDescription([self.data[x.item()] for x in item])
+ elif item.dtype == torch.bool:
+ return ObjDescription(list(compress(self.data, item)))
+
+ return ObjDescription(list(compress(self.data, item)))
+
+ def __len__(self):
+ return len(self.data)
+
+ def __repr__(self):
+ return "ObjDescription({})".format(self.data)
\ No newline at end of file
diff --git a/models/grit_src/grit/data/datasets/__pycache__/grit_coco.cpython-38.pyc b/models/grit_src/grit/data/datasets/__pycache__/grit_coco.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..028e768a57de7b07c2e2ccc75d1d93424cdab4ca
Binary files /dev/null and b/models/grit_src/grit/data/datasets/__pycache__/grit_coco.cpython-38.pyc differ
diff --git a/models/grit_src/grit/data/datasets/__pycache__/object365.cpython-38.pyc b/models/grit_src/grit/data/datasets/__pycache__/object365.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..643a327af0d9748b117fd034ac324a29fe00a60c
Binary files /dev/null and b/models/grit_src/grit/data/datasets/__pycache__/object365.cpython-38.pyc differ
diff --git a/models/grit_src/grit/data/datasets/__pycache__/vg.cpython-38.pyc b/models/grit_src/grit/data/datasets/__pycache__/vg.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..826f8f7d5373054c1eecf273b70efe240aec7e4a
Binary files /dev/null and b/models/grit_src/grit/data/datasets/__pycache__/vg.cpython-38.pyc differ
diff --git a/models/grit_src/grit/data/datasets/grit_coco.py b/models/grit_src/grit/data/datasets/grit_coco.py
new file mode 100644
index 0000000000000000000000000000000000000000..fea81f7dd8ad2c27dac8438753b845ab64cef81e
--- /dev/null
+++ b/models/grit_src/grit/data/datasets/grit_coco.py
@@ -0,0 +1,112 @@
+import logging
+import os
+from fvcore.common.timer import Timer
+from detectron2.structures import BoxMode
+from fvcore.common.file_io import PathManager
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from lvis import LVIS
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["load_GRiTcoco_json", "register_GRiTcoco_instances"]
+
+
+def register_GRiTcoco_instances(name, metadata, json_file, image_root):
+ """
+ """
+ DatasetCatalog.register(name, lambda: load_GRiTcoco_json(
+ json_file, image_root, name))
+ MetadataCatalog.get(name).set(
+ json_file=json_file, image_root=image_root,
+ evaluator_type="coco", **metadata
+ )
+
+
+def get_GRiTcoco_meta():
+ categories = [{'supercategory': 'object', 'id': 1, 'name': 'object'}]
+ categories = sorted(categories, key=lambda x: x["id"])
+ thing_classes = [k["name"] for k in categories]
+ meta = {"thing_classes": thing_classes}
+ return meta
+
+
+def load_GRiTcoco_json(json_file, image_root, dataset_name=None):
+ '''
+ Load COCO class name text for object description for GRiT
+ '''
+
+ json_file = PathManager.get_local_path(json_file)
+
+ timer = Timer()
+ lvis_api = LVIS(json_file)
+ if timer.seconds() > 1:
+ logger.info("Loading {} takes {:.2f} seconds.".format(
+ json_file, timer.seconds()))
+
+ class_names = {}
+ sort_cat = sorted(lvis_api.dataset['categories'], key=lambda x: x['id'])
+ for x in sort_cat:
+ class_names[x['id']] = x['name']
+
+ img_ids = sorted(lvis_api.imgs.keys())
+ imgs = lvis_api.load_imgs(img_ids)
+ anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
+
+ ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+ assert len(set(ann_ids)) == len(ann_ids), \
+ "Annotation ids in '{}' are not unique".format(json_file)
+
+ imgs_anns = list(zip(imgs, anns))
+ logger.info("Loaded {} images in the LVIS v1 format from {}".format(
+ len(imgs_anns), json_file))
+
+ dataset_dicts = []
+
+ for (img_dict, anno_dict_list) in imgs_anns:
+ record = {}
+ if "file_name" in img_dict:
+ file_name = img_dict["file_name"]
+ record["file_name"] = os.path.join(image_root, file_name)
+
+ record["height"] = int(img_dict["height"])
+ record["width"] = int(img_dict["width"])
+ image_id = record["image_id"] = img_dict["id"]
+
+ objs = []
+ for anno in anno_dict_list:
+ assert anno["image_id"] == image_id
+ if anno.get('iscrowd', 0) > 0:
+ continue
+ obj = {"bbox": anno["bbox"], "bbox_mode": BoxMode.XYWH_ABS}
+ obj["category_id"] = 0
+ obj["object_description"] = class_names[anno['category_id']]
+ if 'segmentation' in anno:
+ segm = anno["segmentation"]
+ valid_segm = [poly for poly in segm \
+ if len(poly) % 2 == 0 and len(poly) >= 6]
+ if not len(segm) == len(valid_segm):
+ print('Annotation contains an invalid polygon with < 3 points')
+ assert len(segm) > 0
+ obj["segmentation"] = segm
+ objs.append(obj)
+ record["annotations"] = objs
+ if len(record["annotations"]) == 0:
+ continue
+ record["task"] = "ObjectDet"
+ dataset_dicts.append(record)
+
+ return dataset_dicts
+
+
+_CUSTOM_SPLITS_LVIS = {
+ "GRiT_coco2017_train": ("coco/train2017/", "coco/annotations/instances_train2017.json"),
+}
+
+
+for key, (image_root, json_file) in _CUSTOM_SPLITS_LVIS.items():
+ register_GRiTcoco_instances(
+ key,
+ get_GRiTcoco_meta(),
+ os.path.join("datasets", json_file) if "://" not in json_file else json_file,
+ os.path.join("datasets", image_root),
+ )
\ No newline at end of file
diff --git a/models/grit_src/grit/data/datasets/object365.py b/models/grit_src/grit/data/datasets/object365.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b8cc19da23d8397284b50588ee46e750b5b7552
--- /dev/null
+++ b/models/grit_src/grit/data/datasets/object365.py
@@ -0,0 +1,111 @@
+import logging
+import os
+from fvcore.common.timer import Timer
+from detectron2.structures import BoxMode
+from fvcore.common.file_io import PathManager
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from lvis import LVIS
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["load_o365_json", "register_o365_instances"]
+
+
+def register_o365_instances(name, metadata, json_file, image_root):
+ DatasetCatalog.register(name, lambda: load_o365_json(
+ json_file, image_root, name))
+ MetadataCatalog.get(name).set(
+ json_file=json_file, image_root=image_root,
+ evaluator_type="lvis", **metadata
+ )
+
+
+def get_o365_meta():
+ categories = [{'supercategory': 'object', 'id': 1, 'name': 'object'}]
+ o365_categories = sorted(categories, key=lambda x: x["id"])
+ thing_classes = [k["name"] for k in o365_categories]
+ meta = {"thing_classes": thing_classes}
+ return meta
+
+
+def load_o365_json(json_file, image_root, dataset_name=None):
+ '''
+ Load Object365 class name text for object description for GRiT
+ '''
+
+ json_file = PathManager.get_local_path(json_file)
+
+ timer = Timer()
+ lvis_api = LVIS(json_file)
+ if timer.seconds() > 1:
+ logger.info("Loading {} takes {:.2f} seconds.".format(
+ json_file, timer.seconds()))
+
+ class_names = {}
+ sort_cat = sorted(lvis_api.dataset['categories'], key=lambda x: x['id'])
+ for x in sort_cat:
+ if '/' in x['name']:
+ text = ''
+ for xx in x['name'].split('/'):
+ text += xx
+ text += ' '
+ text = text[:-1]
+ else:
+ text = x['name']
+ class_names[x['id']] = text
+
+ img_ids = sorted(lvis_api.imgs.keys())
+ imgs = lvis_api.load_imgs(img_ids)
+ anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
+
+ ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+ assert len(set(ann_ids)) == len(ann_ids), \
+ "Annotation ids in '{}' are not unique".format(json_file)
+
+ imgs_anns = list(zip(imgs, anns))
+ logger.info("Loaded {} images in the LVIS v1 format from {}".format(
+ len(imgs_anns), json_file))
+
+ dataset_dicts = []
+
+ for (img_dict, anno_dict_list) in imgs_anns:
+ record = {}
+ if "file_name" in img_dict:
+ file_name = img_dict["file_name"]
+ record["file_name"] = os.path.join(image_root, file_name)
+
+ record["height"] = int(img_dict["height"])
+ record["width"] = int(img_dict["width"])
+ image_id = record["image_id"] = img_dict["id"]
+
+ objs = []
+ for anno in anno_dict_list:
+ assert anno["image_id"] == image_id
+ if anno.get('iscrowd', 0) > 0:
+ continue
+ obj = {"bbox": anno["bbox"], "bbox_mode": BoxMode.XYWH_ABS}
+ obj["category_id"] = 0
+ obj["object_description"] = class_names[anno['category_id']]
+
+ objs.append(obj)
+ record["annotations"] = objs
+ if len(record["annotations"]) == 0:
+ continue
+ record["task"] = "ObjectDet"
+ dataset_dicts.append(record)
+
+ return dataset_dicts
+
+
+_CUSTOM_SPLITS_LVIS = {
+ "object365_train": ("object365/images/train/", "object365/annotations/train_v1.json"),
+}
+
+
+for key, (image_root, json_file) in _CUSTOM_SPLITS_LVIS.items():
+ register_o365_instances(
+ key,
+ get_o365_meta(),
+ os.path.join("datasets", json_file) if "://" not in json_file else json_file,
+ os.path.join("datasets", image_root),
+ )
\ No newline at end of file
diff --git a/models/grit_src/grit/data/datasets/vg.py b/models/grit_src/grit/data/datasets/vg.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d47a80d9f88b89ca3064dbc4945b0246162e5d1
--- /dev/null
+++ b/models/grit_src/grit/data/datasets/vg.py
@@ -0,0 +1,98 @@
+import logging
+import os
+from fvcore.common.timer import Timer
+from detectron2.structures import BoxMode
+from fvcore.common.file_io import PathManager
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from lvis import LVIS
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["load_vg_json", "register_vg_instances"]
+
+
+def register_vg_instances(name, metadata, json_file, image_root):
+ """
+ """
+ DatasetCatalog.register(name, lambda: load_vg_json(
+ json_file, image_root, name))
+ MetadataCatalog.get(name).set(
+ json_file=json_file, image_root=image_root,
+ evaluator_type="vg", **metadata
+ )
+
+
+def get_vg_meta():
+ categories = [{'supercategory': 'object', 'id': 1, 'name': 'object'}]
+ vg_categories = sorted(categories, key=lambda x: x["id"])
+ thing_classes = [k["name"] for k in vg_categories]
+ meta = {"thing_classes": thing_classes}
+ return meta
+
+
+def load_vg_json(json_file, image_root, dataset_name=None):
+
+ json_file = PathManager.get_local_path(json_file)
+
+ timer = Timer()
+ lvis_api = LVIS(json_file)
+ if timer.seconds() > 1:
+ logger.info("Loading {} takes {:.2f} seconds.".format(
+ json_file, timer.seconds()))
+
+ img_ids = sorted(lvis_api.imgs.keys())
+ imgs = lvis_api.load_imgs(img_ids)
+ anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
+
+ ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+ assert len(set(ann_ids)) == len(ann_ids), \
+ "Annotation ids in '{}' are not unique".format(json_file)
+
+ imgs_anns = list(zip(imgs, anns))
+ logger.info("Loaded {} images in the LVIS v1 format from {}".format(
+ len(imgs_anns), json_file))
+
+ dataset_dicts = []
+
+ for (img_dict, anno_dict_list) in imgs_anns:
+ record = {}
+ if "file_name" in img_dict:
+ file_name = img_dict["file_name"]
+ record["file_name"] = os.path.join(image_root, file_name)
+
+ record["height"] = int(img_dict["height"])
+ record["width"] = int(img_dict["width"])
+ image_id = record["image_id"] = img_dict["id"]
+
+ objs = []
+ for anno in anno_dict_list:
+ assert anno["image_id"] == image_id
+ if anno.get('iscrowd', 0) > 0:
+ continue
+ obj = {"bbox": anno["bbox"], "bbox_mode": BoxMode.XYWH_ABS}
+ obj["category_id"] = 0
+ obj["object_description"] = anno["caption"]
+
+ objs.append(obj)
+ record["annotations"] = objs
+ if len(record["annotations"]) == 0:
+ continue
+ record["task"] = "DenseCap"
+ dataset_dicts.append(record)
+
+ return dataset_dicts
+
+
+_CUSTOM_SPLITS_LVIS = {
+ "vg_train": ("vg/images", "vg/annotations/train.json"),
+ "vg_test": ("vg/images", "vg/annotations/test.json"),
+}
+
+
+for key, (image_root, json_file) in _CUSTOM_SPLITS_LVIS.items():
+ register_vg_instances(
+ key,
+ get_vg_meta(),
+ os.path.join("datasets", json_file) if "://" not in json_file else json_file,
+ os.path.join("datasets", image_root),
+ )
\ No newline at end of file
diff --git a/models/grit_src/grit/data/transforms/__pycache__/custom_augmentation_impl.cpython-38.pyc b/models/grit_src/grit/data/transforms/__pycache__/custom_augmentation_impl.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..352f7cc7ca9734a665a8a630c3f33234e7b48d17
Binary files /dev/null and b/models/grit_src/grit/data/transforms/__pycache__/custom_augmentation_impl.cpython-38.pyc differ
diff --git a/models/grit_src/grit/data/transforms/__pycache__/custom_transform.cpython-38.pyc b/models/grit_src/grit/data/transforms/__pycache__/custom_transform.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5a194459b13bb16d30360ae1883d7e2d01d4a2f1
Binary files /dev/null and b/models/grit_src/grit/data/transforms/__pycache__/custom_transform.cpython-38.pyc differ
diff --git a/models/grit_src/grit/data/transforms/custom_augmentation_impl.py b/models/grit_src/grit/data/transforms/custom_augmentation_impl.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b9637f3ad41e3ba513636219e49371296d9ab9f
--- /dev/null
+++ b/models/grit_src/grit/data/transforms/custom_augmentation_impl.py
@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Part of the code is from https://github.com/rwightman/efficientdet-pytorch/blob/master/effdet/data/transforms.py
+# Modified by Xingyi Zhou
+# The original code is under Apache-2.0 License
+import numpy as np
+from PIL import Image
+
+from detectron2.data.transforms.augmentation import Augmentation
+from .custom_transform import EfficientDetResizeCropTransform
+
+__all__ = [
+ "EfficientDetResizeCrop",
+]
+
+
+class EfficientDetResizeCrop(Augmentation):
+ """
+ Scale the shorter edge to the given size, with a limit of `max_size` on the longer edge.
+ If `max_size` is reached, then downscale so that the longer edge does not exceed max_size.
+ """
+
+ def __init__(
+ self, size, scale, interp=Image.BILINEAR
+ ):
+ """
+ """
+ super().__init__()
+ self.target_size = (size, size)
+ self.scale = scale
+ self.interp = interp
+
+ def get_transform(self, img):
+ # Select a random scale factor.
+ scale_factor = np.random.uniform(*self.scale)
+ scaled_target_height = scale_factor * self.target_size[0]
+ scaled_target_width = scale_factor * self.target_size[1]
+ # Recompute the accurate scale_factor using rounded scaled image size.
+ width, height = img.shape[1], img.shape[0]
+ img_scale_y = scaled_target_height / height
+ img_scale_x = scaled_target_width / width
+ img_scale = min(img_scale_y, img_scale_x)
+
+ # Select non-zero random offset (x, y) if scaled image is larger than target size
+ scaled_h = int(height * img_scale)
+ scaled_w = int(width * img_scale)
+ offset_y = scaled_h - self.target_size[0]
+ offset_x = scaled_w - self.target_size[1]
+ offset_y = int(max(0.0, float(offset_y)) * np.random.uniform(0, 1))
+ offset_x = int(max(0.0, float(offset_x)) * np.random.uniform(0, 1))
+ return EfficientDetResizeCropTransform(
+ scaled_h, scaled_w, offset_y, offset_x, img_scale, self.target_size, self.interp)
diff --git a/models/grit_src/grit/data/transforms/custom_transform.py b/models/grit_src/grit/data/transforms/custom_transform.py
new file mode 100644
index 0000000000000000000000000000000000000000..423063a4ea14fe92caaed7efc69d8596a597485e
--- /dev/null
+++ b/models/grit_src/grit/data/transforms/custom_transform.py
@@ -0,0 +1,115 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Part of the code is from https://github.com/rwightman/efficientdet-pytorch/blob/master/effdet/data/transforms.py
+# Modified by Xingyi Zhou
+# The original code is under Apache-2.0 License
+import numpy as np
+import torch
+import torch.nn.functional as F
+from fvcore.transforms.transform import (
+ CropTransform,
+ HFlipTransform,
+ NoOpTransform,
+ Transform,
+ TransformList,
+)
+from PIL import Image
+
+try:
+ import cv2 # noqa
+except ImportError:
+ # OpenCV is an optional dependency at the moment
+ pass
+
+__all__ = [
+ "EfficientDetResizeCropTransform",
+]
+
+
+class EfficientDetResizeCropTransform(Transform):
+ """
+ """
+
+ def __init__(self, scaled_h, scaled_w, offset_y, offset_x, img_scale, \
+ target_size, interp=None):
+ """
+ Args:
+ h, w (int): original image size
+ new_h, new_w (int): new image size
+ interp: PIL interpolation methods, defaults to bilinear.
+ """
+ # TODO decide on PIL vs opencv
+ super().__init__()
+ if interp is None:
+ interp = Image.BILINEAR
+ self._set_attributes(locals())
+
+ def apply_image(self, img, interp=None):
+ assert len(img.shape) <= 4
+
+ if img.dtype == np.uint8:
+ pil_image = Image.fromarray(img)
+ interp_method = interp if interp is not None else self.interp
+ pil_image = pil_image.resize((self.scaled_w, self.scaled_h), interp_method)
+ ret = np.asarray(pil_image)
+ right = min(self.scaled_w, self.offset_x + self.target_size[1])
+ lower = min(self.scaled_h, self.offset_y + self.target_size[0])
+ if len(ret.shape) <= 3:
+ ret = ret[self.offset_y: lower, self.offset_x: right]
+ else:
+ ret = ret[..., self.offset_y: lower, self.offset_x: right, :]
+ else:
+ # PIL only supports uint8
+ img = torch.from_numpy(img)
+ shape = list(img.shape)
+ shape_4d = shape[:2] + [1] * (4 - len(shape)) + shape[2:]
+ img = img.view(shape_4d).permute(2, 3, 0, 1) # hw(c) -> nchw
+ _PIL_RESIZE_TO_INTERPOLATE_MODE = {Image.BILINEAR: "bilinear", Image.BICUBIC: "bicubic"}
+ mode = _PIL_RESIZE_TO_INTERPOLATE_MODE[self.interp]
+ img = F.interpolate(img, (self.scaled_h, self.scaled_w), mode=mode, align_corners=False)
+ shape[:2] = (self.scaled_h, self.scaled_w)
+ ret = img.permute(2, 3, 0, 1).view(shape).numpy() # nchw -> hw(c)
+ right = min(self.scaled_w, self.offset_x + self.target_size[1])
+ lower = min(self.scaled_h, self.offset_y + self.target_size[0])
+ if len(ret.shape) <= 3:
+ ret = ret[self.offset_y: lower, self.offset_x: right]
+ else:
+ ret = ret[..., self.offset_y: lower, self.offset_x: right, :]
+ return ret
+
+
+ def apply_coords(self, coords):
+ coords[:, 0] = coords[:, 0] * self.img_scale
+ coords[:, 1] = coords[:, 1] * self.img_scale
+ coords[:, 0] -= self.offset_x
+ coords[:, 1] -= self.offset_y
+ return coords
+
+
+ def apply_segmentation(self, segmentation):
+ segmentation = self.apply_image(segmentation, interp=Image.NEAREST)
+ return segmentation
+
+
+ def inverse(self):
+ raise NotImplementedError
+
+
+ def inverse_apply_coords(self, coords):
+ coords[:, 0] += self.offset_x
+ coords[:, 1] += self.offset_y
+ coords[:, 0] = coords[:, 0] / self.img_scale
+ coords[:, 1] = coords[:, 1] / self.img_scale
+ return coords
+
+
+ def inverse_apply_box(self, box: np.ndarray) -> np.ndarray:
+ """
+ """
+ idxs = np.array([(0, 1), (2, 1), (0, 3), (2, 3)]).flatten()
+ coords = np.asarray(box).reshape(-1, 4)[:, idxs].reshape(-1, 2)
+ coords = self.inverse_apply_coords(coords).reshape((-1, 4, 2))
+ minxy = coords.min(axis=1)
+ maxxy = coords.max(axis=1)
+ trans_boxes = np.concatenate((minxy, maxxy), axis=1)
+ return trans_boxes
\ No newline at end of file
diff --git a/models/grit_src/grit/evaluation/eval.py b/models/grit_src/grit/evaluation/eval.py
new file mode 100644
index 0000000000000000000000000000000000000000..951a0920ec3d93703245562d4f76ec597e672ad9
--- /dev/null
+++ b/models/grit_src/grit/evaluation/eval.py
@@ -0,0 +1,156 @@
+import itertools
+import json
+import os
+from detectron2.structures import Boxes, BoxMode, pairwise_iou
+from detectron2.utils.file_io import PathManager
+import numpy as np
+import pycocotools.mask as mask_util
+from detectron2.evaluation.coco_evaluation import COCOEvaluator
+from detectron2.evaluation.coco_evaluation import _evaluate_predictions_on_coco
+
+
+class GRiTCOCOEvaluator(COCOEvaluator):
+ def process(self, inputs, outputs):
+ for input, output in zip(inputs, outputs):
+ prediction = {"image_id": input["image_id"]}
+
+ if "instances" in output:
+ instances = output["instances"].to(self._cpu_device)
+ prediction["instances"] = instances_to_coco_json(instances, input["image_id"])
+
+ if len(prediction) > 1:
+ self._predictions.append(prediction)
+
+ def _eval_predictions(self, predictions, img_ids=None):
+ self._logger.info("Preparing results for COCO format ...")
+ coco_results = list(itertools.chain(*[x["instances"] for x in predictions]))
+ tasks = self._tasks or self._tasks_from_predictions(coco_results)
+
+ if self._output_dir:
+ file_path = os.path.join(self._output_dir, "coco_instances_results.json")
+ self._logger.info("Saving results to {}".format(file_path))
+ with PathManager.open(file_path, "w") as f:
+ f.write(json.dumps(coco_results))
+ f.flush()
+
+ if not self._do_evaluation:
+ self._logger.info("Annotations are not available for evaluation.")
+ return
+
+ self._logger.info(
+ "Evaluating predictions with {} COCO API...".format(
+ "unofficial" if self._use_fast_impl else "official"
+ )
+ )
+
+ coco_results = self.convert_classname_to_id(coco_results)
+
+ for task in sorted(tasks):
+ assert task in {"bbox", "segm", "keypoints"}, f"Got unknown task: {task}!"
+ coco_eval = (
+ _evaluate_predictions_on_coco(
+ self._coco_api,
+ coco_results,
+ task,
+ kpt_oks_sigmas=self._kpt_oks_sigmas,
+ use_fast_impl=self._use_fast_impl,
+ img_ids=img_ids,
+ max_dets_per_image=self._max_dets_per_image,
+ )
+ if len(coco_results) > 0
+ else None # cocoapi does not handle empty results very well
+ )
+
+ res = self._derive_coco_results(
+ coco_eval, task, class_names=self._metadata.get("thing_classes")
+ )
+ self._results[task] = res
+
+ def convert_classname_to_id(self, results):
+ outputs = []
+ class_name_to_id = {}
+ categories = sorted(self._coco_api.dataset['categories'], key=lambda x: x['id'])
+
+ for cat in categories:
+ class_name_to_id[cat['name']] = cat['id']
+
+ for pred in results:
+ if pred['object_descriptions'] in class_name_to_id:
+ pred['category_id'] = class_name_to_id[pred['object_descriptions']]
+ del pred['object_descriptions']
+ outputs.append(pred)
+
+ return outputs
+
+
+class GRiTVGEvaluator(COCOEvaluator):
+ def process(self, inputs, outputs):
+ for input, output in zip(inputs, outputs):
+ assert input["image_id"] == int(input['file_name'].split('/')[-1].split('.')[0])
+ prediction = {"image_id": input["image_id"]}
+
+ if "instances" in output:
+ instances = output["instances"].to(self._cpu_device)
+ prediction["instances"] = instances_to_coco_json(instances, input["image_id"], output_logits=True)
+ h = input['height']
+ w = input['width']
+ scale = 720.0 / max(h, w)
+ scaled_inst = []
+ for inst in prediction["instances"]:
+ inst['bbox'][0] = inst['bbox'][0] * scale
+ inst['bbox'][1] = inst['bbox'][1] * scale
+ inst['bbox'][2] = inst['bbox'][2] * scale
+ inst['bbox'][3] = inst['bbox'][3] * scale
+ scaled_inst.append(inst)
+ if len(scaled_inst) > 0:
+ prediction["instances"] = scaled_inst
+ if len(prediction) > 1:
+ self._predictions.append(prediction)
+
+ def _eval_predictions(self, predictions, img_ids=None):
+ '''
+ This is only for saving the results to json file
+ '''
+ self._logger.info("Preparing results for COCO format ...")
+ coco_results = list(itertools.chain(*[x["instances"] for x in predictions]))
+
+ if self._output_dir:
+ file_path = os.path.join(self._output_dir, "vg_instances_results.json")
+ self._logger.info("Saving results to {}".format(file_path))
+ with PathManager.open(file_path, "w") as f:
+ f.write(json.dumps(coco_results))
+ f.flush()
+
+
+def instances_to_coco_json(instances, img_id, output_logits=False):
+ """
+ Add object_descriptions and logit (if applicable) to
+ detectron2's instances_to_coco_json
+ """
+ num_instance = len(instances)
+ if num_instance == 0:
+ return []
+
+ boxes = instances.pred_boxes.tensor.numpy()
+ boxes = BoxMode.convert(boxes, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+ boxes = boxes.tolist()
+ scores = instances.scores.tolist()
+ classes = instances.pred_classes.tolist()
+ object_descriptions = instances.pred_object_descriptions.data
+ if output_logits:
+ logits = instances.logits.tolist()
+
+ results = []
+ for k in range(num_instance):
+ result = {
+ "image_id": img_id,
+ "category_id": classes[k],
+ "bbox": boxes[k],
+ "score": scores[k],
+ 'object_descriptions': object_descriptions[k],
+ }
+ if output_logits:
+ result["logit"] = logits[k]
+
+ results.append(result)
+ return results
\ No newline at end of file
diff --git a/models/grit_src/grit/modeling/__pycache__/soft_nms.cpython-38.pyc b/models/grit_src/grit/modeling/__pycache__/soft_nms.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..49cddaf6500323a7989b833e1551006ff02bc73b
Binary files /dev/null and b/models/grit_src/grit/modeling/__pycache__/soft_nms.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/backbone/__pycache__/utils.cpython-38.pyc b/models/grit_src/grit/modeling/backbone/__pycache__/utils.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..58de6f675ddaea4e32d4f131d40917b28b910546
Binary files /dev/null and b/models/grit_src/grit/modeling/backbone/__pycache__/utils.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/backbone/__pycache__/vit.cpython-38.pyc b/models/grit_src/grit/modeling/backbone/__pycache__/vit.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2bb5bf0d51fe6dd3de6571884ff1b2ec02653f8b
Binary files /dev/null and b/models/grit_src/grit/modeling/backbone/__pycache__/vit.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/backbone/utils.py b/models/grit_src/grit/modeling/backbone/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e71db21f1223c87cceeb422a70888f7bac42bb18
--- /dev/null
+++ b/models/grit_src/grit/modeling/backbone/utils.py
@@ -0,0 +1,186 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# This code is from https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/utils.py
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = [
+ "window_partition",
+ "window_unpartition",
+ "add_decomposed_rel_pos",
+ "get_abs_pos",
+ "PatchEmbed",
+]
+
+def window_partition(x, window_size):
+ """
+ Partition into non-overlapping windows with padding if needed.
+ Args:
+ x (tensor): input tokens with [B, H, W, C].
+ window_size (int): window size.
+
+ Returns:
+ windows: windows after partition with [B * num_windows, window_size, window_size, C].
+ (Hp, Wp): padded height and width before partition
+ """
+ B, H, W, C = x.shape
+
+ pad_h = (window_size - H % window_size) % window_size
+ pad_w = (window_size - W % window_size) % window_size
+ if pad_h > 0 or pad_w > 0:
+ x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+ Hp, Wp = H + pad_h, W + pad_w
+
+ x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+ return windows, (Hp, Wp)
+
+
+def window_unpartition(windows, window_size, pad_hw, hw):
+ """
+ Window unpartition into original sequences and removing padding.
+ Args:
+ x (tensor): input tokens with [B * num_windows, window_size, window_size, C].
+ window_size (int): window size.
+ pad_hw (Tuple): padded height and width (Hp, Wp).
+ hw (Tuple): original height and width (H, W) before padding.
+
+ Returns:
+ x: unpartitioned sequences with [B, H, W, C].
+ """
+ Hp, Wp = pad_hw
+ H, W = hw
+ B = windows.shape[0] // (Hp * Wp // window_size // window_size)
+ x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)
+
+ if Hp > H or Wp > W:
+ x = x[:, :H, :W, :].contiguous()
+ return x
+
+
+def get_rel_pos(q_size, k_size, rel_pos):
+ """
+ Get relative positional embeddings according to the relative positions of
+ query and key sizes.
+ Args:
+ q_size (int): size of query q.
+ k_size (int): size of key k.
+ rel_pos (Tensor): relative position embeddings (L, C).
+
+ Returns:
+ Extracted positional embeddings according to relative positions.
+ """
+ max_rel_dist = int(2 * max(q_size, k_size) - 1)
+ # Interpolate rel pos if needed.
+ if rel_pos.shape[0] != max_rel_dist:
+ # Interpolate rel pos.
+ rel_pos_resized = F.interpolate(
+ rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
+ size=max_rel_dist,
+ mode="linear",
+ )
+ rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
+ else:
+ rel_pos_resized = rel_pos
+
+ # Scale the coords with short length if shapes for q and k are different.
+ q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
+ k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
+ relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
+
+ return rel_pos_resized[relative_coords.long()]
+
+
+def add_decomposed_rel_pos(attn, q, rel_pos_h, rel_pos_w, q_size, k_size):
+ """
+ Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+ https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950
+ Args:
+ attn (Tensor): attention map.
+ q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
+ rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.
+ rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.
+ q_size (Tuple): spatial sequence size of query q with (q_h, q_w).
+ k_size (Tuple): spatial sequence size of key k with (k_h, k_w).
+
+ Returns:
+ attn (Tensor): attention map with added relative positional embeddings.
+ """
+ q_h, q_w = q_size
+ k_h, k_w = k_size
+ Rh = get_rel_pos(q_h, k_h, rel_pos_h)
+ Rw = get_rel_pos(q_w, k_w, rel_pos_w)
+
+ B, _, dim = q.shape
+ r_q = q.reshape(B, q_h, q_w, dim)
+ rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)
+ rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)
+
+ attn = (
+ attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
+ ).view(B, q_h * q_w, k_h * k_w)
+
+ return attn
+
+
+def get_abs_pos(abs_pos, has_cls_token, hw):
+ """
+ Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token
+ dimension for the original embeddings.
+ Args:
+ abs_pos (Tensor): absolute positional embeddings with (1, num_position, C).
+ has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token.
+ hw (Tuple): size of input image tokens.
+
+ Returns:
+ Absolute positional embeddings after processing with shape (1, H, W, C)
+ """
+ h, w = hw
+ if has_cls_token:
+ abs_pos = abs_pos[:, 1:]
+ xy_num = abs_pos.shape[1]
+ size = int(math.sqrt(xy_num))
+ assert size * size == xy_num
+
+ if size != h or size != w:
+ new_abs_pos = F.interpolate(
+ abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2),
+ size=(h, w),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ return new_abs_pos.permute(0, 2, 3, 1)
+ else:
+ return abs_pos.reshape(1, h, w, -1)
+
+
+class PatchEmbed(nn.Module):
+ """
+ Image to Patch Embedding.
+ """
+
+ def __init__(
+ self, kernel_size=(16, 16), stride=(16, 16), padding=(0, 0), in_chans=3, embed_dim=768
+ ):
+ """
+ Args:
+ kernel_size (Tuple): kernel size of the projection layer.
+ stride (Tuple): stride of the projection layer.
+ padding (Tuple): padding size of the projection layer.
+ in_chans (int): Number of input image channels.
+ embed_dim (int): embed_dim (int): Patch embedding dimension.
+ """
+ super().__init__()
+
+ self.proj = nn.Conv2d(
+ in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
+ )
+
+ def forward(self, x):
+ x = self.proj(x)
+ # B C H W -> B H W C
+ x = x.permute(0, 2, 3, 1)
+ return x
diff --git a/models/grit_src/grit/modeling/backbone/vit.py b/models/grit_src/grit/modeling/backbone/vit.py
new file mode 100644
index 0000000000000000000000000000000000000000..36d1207f1b5a59030b77984e424a4b1e8b3c761a
--- /dev/null
+++ b/models/grit_src/grit/modeling/backbone/vit.py
@@ -0,0 +1,538 @@
+# Modified by Jialian Wu from https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py
+import logging
+import math
+import fvcore.nn.weight_init as weight_init
+import torch
+import torch.nn as nn
+from functools import partial
+
+from detectron2.layers import CNNBlockBase, Conv2d, get_norm
+from detectron2.modeling.backbone.build import BACKBONE_REGISTRY
+from detectron2.layers import ShapeSpec
+from centernet.modeling.backbone.fpn_p5 import LastLevelP6P7_P5
+
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, Mlp, trunc_normal_
+
+from detectron2.modeling.backbone.backbone import Backbone
+from .utils import (
+ PatchEmbed,
+ add_decomposed_rel_pos,
+ get_abs_pos,
+ window_partition,
+ window_unpartition,
+)
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = ["ViT"]
+
+
+class Attention(nn.Module):
+ """Multi-head Attention block with relative position embeddings."""
+
+ def __init__(
+ self,
+ dim,
+ num_heads=8,
+ qkv_bias=True,
+ use_rel_pos=False,
+ rel_pos_zero_init=True,
+ input_size=None,
+ ):
+ """
+ Args:
+ dim (int): Number of input channels.
+ num_heads (int): Number of attention heads.
+ qkv_bias (bool: If True, add a learnable bias to query, key, value.
+ rel_pos (bool): If True, add relative positional embeddings to the attention map.
+ rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+ input_size (int or None): Input resolution for calculating the relative positional
+ parameter size.
+ """
+ super().__init__()
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ self.scale = head_dim**-0.5
+
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+ self.proj = nn.Linear(dim, dim)
+
+ self.use_rel_pos = use_rel_pos
+ if self.use_rel_pos:
+ # initialize relative positional embeddings
+ self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
+ self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
+
+ if not rel_pos_zero_init:
+ trunc_normal_(self.rel_pos_h, std=0.02)
+ trunc_normal_(self.rel_pos_w, std=0.02)
+
+ def forward(self, x):
+ B, H, W, _ = x.shape
+ # qkv with shape (3, B, nHead, H * W, C)
+ qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+ # q, k, v with shape (B * nHead, H * W, C)
+ q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
+
+ attn = (q * self.scale) @ k.transpose(-2, -1)
+
+ if self.use_rel_pos:
+ attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))
+
+ attn = attn.softmax(dim=-1)
+ x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1)
+ x = self.proj(x)
+
+ return x
+
+
+class ResBottleneckBlock(CNNBlockBase):
+ """
+ The standard bottleneck residual block without the last activation layer.
+ It contains 3 conv layers with kernels 1x1, 3x3, 1x1.
+ """
+
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ bottleneck_channels,
+ norm="LN",
+ act_layer=nn.GELU,
+ ):
+ """
+ Args:
+ in_channels (int): Number of input channels.
+ out_channels (int): Number of output channels.
+ bottleneck_channels (int): number of output channels for the 3x3
+ "bottleneck" conv layers.
+ norm (str or callable): normalization for all conv layers.
+ See :func:`layers.get_norm` for supported format.
+ act_layer (callable): activation for all conv layers.
+ """
+ super().__init__(in_channels, out_channels, 1)
+
+ self.conv1 = Conv2d(in_channels, bottleneck_channels, 1, bias=False)
+ self.norm1 = get_norm(norm, bottleneck_channels)
+ self.act1 = act_layer()
+
+ self.conv2 = Conv2d(
+ bottleneck_channels,
+ bottleneck_channels,
+ 3,
+ padding=1,
+ bias=False,
+ )
+ self.norm2 = get_norm(norm, bottleneck_channels)
+ self.act2 = act_layer()
+
+ self.conv3 = Conv2d(bottleneck_channels, out_channels, 1, bias=False)
+ self.norm3 = get_norm(norm, out_channels)
+
+ for layer in [self.conv1, self.conv2, self.conv3]:
+ weight_init.c2_msra_fill(layer)
+ for layer in [self.norm1, self.norm2]:
+ layer.weight.data.fill_(1.0)
+ layer.bias.data.zero_()
+ # zero init last norm layer.
+ self.norm3.weight.data.zero_()
+ self.norm3.bias.data.zero_()
+
+ def forward(self, x):
+ out = x
+ for layer in self.children():
+ out = layer(out)
+
+ out = x + out
+ return out
+
+
+class Block(nn.Module):
+ """Transformer blocks with support of window attention and residual propagation blocks"""
+
+ def __init__(
+ self,
+ dim,
+ num_heads,
+ mlp_ratio=4.0,
+ qkv_bias=True,
+ drop_path=0.0,
+ norm_layer=nn.LayerNorm,
+ act_layer=nn.GELU,
+ use_rel_pos=False,
+ rel_pos_zero_init=True,
+ window_size=0,
+ use_residual_block=False,
+ input_size=None,
+ ):
+ """
+ Args:
+ dim (int): Number of input channels.
+ num_heads (int): Number of attention heads in each ViT block.
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+ qkv_bias (bool): If True, add a learnable bias to query, key, value.
+ drop_path (float): Stochastic depth rate.
+ norm_layer (nn.Module): Normalization layer.
+ act_layer (nn.Module): Activation layer.
+ use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+ rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+ window_size (int): Window size for window attention blocks. If it equals 0, then not
+ use window attention.
+ use_residual_block (bool): If True, use a residual block after the MLP block.
+ input_size (int or None): Input resolution for calculating the relative positional
+ parameter size.
+ """
+ super().__init__()
+ self.norm1 = norm_layer(dim)
+ self.attn = Attention(
+ dim,
+ num_heads=num_heads,
+ qkv_bias=qkv_bias,
+ use_rel_pos=use_rel_pos,
+ rel_pos_zero_init=rel_pos_zero_init,
+ input_size=input_size if window_size == 0 else (window_size, window_size),
+ )
+
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+ self.norm2 = norm_layer(dim)
+ self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer)
+
+ self.window_size = window_size
+
+ self.use_residual_block = use_residual_block
+ if use_residual_block:
+ # Use a residual block with bottleneck channel as dim // 2
+ self.residual = ResBottleneckBlock(
+ in_channels=dim,
+ out_channels=dim,
+ bottleneck_channels=dim // 2,
+ norm="LN",
+ act_layer=act_layer,
+ )
+
+ def forward(self, x):
+ shortcut = x
+ x = self.norm1(x)
+ # Window partition
+ if self.window_size > 0:
+ H, W = x.shape[1], x.shape[2]
+ x, pad_hw = window_partition(x, self.window_size)
+
+ x = self.attn(x)
+ # Reverse window partition
+ if self.window_size > 0:
+ x = window_unpartition(x, self.window_size, pad_hw, (H, W))
+
+ x = shortcut + self.drop_path(x)
+ x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+ if self.use_residual_block:
+ x = self.residual(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
+
+ return x
+
+
+class ViT(Backbone):
+ """
+ This module implements Vision Transformer (ViT) backbone in :paper:`vitdet`.
+ "Exploring Plain Vision Transformer Backbones for Object Detection",
+ https://arxiv.org/abs/2203.16527
+ """
+
+ def __init__(
+ self,
+ img_size=1024,
+ patch_size=16,
+ in_chans=3,
+ embed_dim=768,
+ depth=12,
+ num_heads=12,
+ mlp_ratio=4.0,
+ qkv_bias=True,
+ drop_path_rate=0.0,
+ norm_layer=nn.LayerNorm,
+ act_layer=nn.GELU,
+ use_abs_pos=True,
+ use_rel_pos=False,
+ rel_pos_zero_init=True,
+ window_size=0,
+ window_block_indexes=(),
+ residual_block_indexes=(),
+ use_act_checkpoint=True,
+ pretrain_img_size=224,
+ pretrain_use_cls_token=True,
+ out_feature="last_feat",
+ ):
+ """
+ Args:
+ img_size (int): Input image size.
+ patch_size (int): Patch size.
+ in_chans (int): Number of input image channels.
+ embed_dim (int): Patch embedding dimension.
+ depth (int): Depth of ViT.
+ num_heads (int): Number of attention heads in each ViT block.
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+ qkv_bias (bool): If True, add a learnable bias to query, key, value.
+ drop_path_rate (float): Stochastic depth rate.
+ norm_layer (nn.Module): Normalization layer.
+ act_layer (nn.Module): Activation layer.
+ use_abs_pos (bool): If True, use absolute positional embeddings.
+ use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+ rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+ window_size (int): Window size for window attention blocks.
+ window_block_indexes (list): Indexes for blocks using window attention.
+ residual_block_indexes (list): Indexes for blocks using conv propagation.
+ use_act_checkpoint (bool): If True, use activation checkpointing.
+ pretrain_img_size (int): input image size for pretraining models.
+ pretrain_use_cls_token (bool): If True, pretrainig models use class token.
+ out_feature (str): name of the feature from the last block.
+ """
+ super().__init__()
+ self.pretrain_use_cls_token = pretrain_use_cls_token
+ self.use_act_checkpoint = use_act_checkpoint
+
+ self.patch_embed = PatchEmbed(
+ kernel_size=(patch_size, patch_size),
+ stride=(patch_size, patch_size),
+ in_chans=in_chans,
+ embed_dim=embed_dim,
+ )
+
+ if use_abs_pos:
+ # Initialize absolute positional embedding with pretrain image size.
+ num_patches = (pretrain_img_size // patch_size) * (pretrain_img_size // patch_size)
+ num_positions = (num_patches + 1) if pretrain_use_cls_token else num_patches
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_positions, embed_dim))
+ else:
+ self.pos_embed = None
+
+ # stochastic depth decay rule
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+
+ self.blocks = nn.ModuleList()
+ for i in range(depth):
+ block = Block(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ drop_path=dpr[i],
+ norm_layer=norm_layer,
+ act_layer=act_layer,
+ use_rel_pos=use_rel_pos,
+ rel_pos_zero_init=rel_pos_zero_init,
+ window_size=window_size if i in window_block_indexes else 0,
+ use_residual_block=i in residual_block_indexes,
+ input_size=(img_size // patch_size, img_size // patch_size),
+ )
+ self.blocks.append(block)
+
+ self._out_feature_channels = {out_feature: embed_dim}
+ self._out_feature_strides = {out_feature: patch_size}
+ self._out_features = [out_feature]
+
+ if self.pos_embed is not None:
+ trunc_normal_(self.pos_embed, std=0.02)
+
+ self.apply(self._init_weights)
+
+ def _init_weights(self, m):
+ if isinstance(m, nn.Linear):
+ trunc_normal_(m.weight, std=0.02)
+ if isinstance(m, nn.Linear) and m.bias is not None:
+ nn.init.constant_(m.bias, 0)
+ elif isinstance(m, nn.LayerNorm):
+ nn.init.constant_(m.bias, 0)
+ nn.init.constant_(m.weight, 1.0)
+
+ def forward(self, x):
+ x = self.patch_embed(x)
+ if self.pos_embed is not None:
+ x = x + get_abs_pos(
+ self.pos_embed, self.pretrain_use_cls_token, (x.shape[1], x.shape[2])
+ )
+
+ for blk in self.blocks:
+ if self.use_act_checkpoint:
+ x = checkpoint.checkpoint(blk, x)
+ else:
+ x = blk(x)
+
+ return x.permute(0, 3, 1, 2)
+
+
+class ViT_FPN(Backbone):
+ def __init__(self, bottom_up=None, top_block=None, out_channels=None, strides=None, vit_out_dim=None):
+ super(ViT_FPN, self).__init__()
+ assert isinstance(bottom_up, Backbone)
+ self.bottom_up = bottom_up
+ self.top_block = top_block
+
+ self._out_feature_strides = {"p{}".format(int(math.log2(s))): s for s in strides}
+ self._out_features = list(self._out_feature_strides.keys())
+ self._out_feature_channels = {k: out_channels for k in self._out_features}
+ self._size_divisibility = strides[2]
+
+ self.maxpool = nn.MaxPool2d(2, stride=2)
+ self.fpn_stride_16_8 = nn.ConvTranspose2d(vit_out_dim, vit_out_dim, 2, stride=2, bias=False)
+ self.fpn_stride8_conv1 = nn.Conv2d(in_channels=vit_out_dim, out_channels=out_channels, kernel_size=1, bias=False)
+ self.fpn_stride8_norm1 = nn.LayerNorm(out_channels)
+ self.fpn_stride8_conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+ self.fpn_stride8_norm2 = nn.LayerNorm(out_channels)
+
+ self.fpn_stride16_conv1 = nn.Conv2d(in_channels=vit_out_dim, out_channels=out_channels, kernel_size=1, bias=False)
+ self.fpn_stride16_norm1 = nn.LayerNorm(out_channels)
+ self.fpn_stride16_conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+ self.fpn_stride16_norm2 = nn.LayerNorm(out_channels)
+
+ self.fpn_stride32_conv1 = nn.Conv2d(in_channels=vit_out_dim, out_channels=out_channels, kernel_size=1, bias=False)
+ self.fpn_stride32_norm1 = nn.LayerNorm(out_channels)
+ self.fpn_stride32_conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+ self.fpn_stride32_norm2 = nn.LayerNorm(out_channels)
+
+ def forward(self, x):
+ vit_output_featuremap = self.bottom_up(x)
+
+ stride8_feature = self.fpn_stride_16_8(vit_output_featuremap)
+ stride8_feature = self.fpn_stride8_norm1(self.fpn_stride8_conv1(stride8_feature)
+ .permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+ stride8_feature = self.fpn_stride8_norm2(self.fpn_stride8_conv2(stride8_feature)
+ .permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+ stride32_feature = self.maxpool(vit_output_featuremap)
+ stride32_feature = self.fpn_stride32_norm1(self.fpn_stride32_conv1(stride32_feature)
+ .permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+ stride32_feature = self.fpn_stride32_norm2(self.fpn_stride32_conv2(stride32_feature)
+ .permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+ stride16_feature = self.fpn_stride16_norm1(self.fpn_stride16_conv1(vit_output_featuremap).
+ permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+ stride16_feature = self.fpn_stride16_norm2(self.fpn_stride16_conv2(stride16_feature)
+ .permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+ results = [stride8_feature, stride16_feature, stride32_feature]
+
+ results.extend(self.top_block(stride32_feature))
+
+ assert len(self._out_features) == len(results)
+ fpn_out = {f: res for f, res in zip(self._out_features, results)}
+
+ return fpn_out
+ @property
+ def size_divisibility(self):
+ return self._size_divisibility
+
+ def output_shape(self):
+ return {
+ name: ShapeSpec(
+ channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
+ )
+ for name in self._out_features
+ }
+
+
+@BACKBONE_REGISTRY.register()
+def build_vit_fpn_backbone(cfg, input_shape: ShapeSpec):
+ embed_dim = 768
+ vit_out_dim = embed_dim
+ bottom_up = ViT( # Single-scale ViT backbone
+ img_size=1024,
+ patch_size=16,
+ embed_dim=embed_dim,
+ depth=12,
+ num_heads=12,
+ drop_path_rate=0.1,
+ window_size=14,
+ mlp_ratio=4,
+ qkv_bias=True,
+ norm_layer=partial(nn.LayerNorm, eps=1e-6),
+ window_block_indexes=[
+ # 2, 5, 8 11 for global attention
+ 0,
+ 1,
+ 3,
+ 4,
+ 6,
+ 7,
+ 9,
+ 10,
+ ],
+ residual_block_indexes=[],
+ use_act_checkpoint=cfg.USE_ACT_CHECKPOINT,
+ use_rel_pos=True,
+ out_feature="last_feat",)
+
+ out_channels = cfg.MODEL.FPN.OUT_CHANNELS
+ assert out_channels == 256 or out_channels == 768 or out_channels == 1024
+ backbone = ViT_FPN(bottom_up=bottom_up,
+ top_block=LastLevelP6P7_P5(out_channels, out_channels),
+ out_channels=out_channels,
+ strides=[8, 16, 32, 64, 128],
+ vit_out_dim=vit_out_dim)
+ return backbone
+
+
+@BACKBONE_REGISTRY.register()
+def build_vit_fpn_backbone_large(cfg, input_shape: ShapeSpec):
+ window_block_indexes = (list(range(0, 5)) + list(range(6, 11)) + list(range(12, 17)) + list(range(18, 23)))
+ embed_dim = 1024
+ vit_out_dim = embed_dim
+ bottom_up = ViT( # Single-scale ViT backbone
+ img_size=1024,
+ patch_size=16,
+ embed_dim=embed_dim,
+ depth=24,
+ num_heads=16,
+ drop_path_rate=0.4,
+ window_size=14,
+ mlp_ratio=4,
+ qkv_bias=True,
+ norm_layer=partial(nn.LayerNorm, eps=1e-6),
+ window_block_indexes=window_block_indexes,
+ residual_block_indexes=[],
+ use_act_checkpoint=cfg.USE_ACT_CHECKPOINT,
+ use_rel_pos=True,
+ out_feature="last_feat",)
+
+ out_channels = cfg.MODEL.FPN.OUT_CHANNELS
+ assert out_channels == 256 or out_channels == 768 or out_channels == 1024
+ backbone = ViT_FPN(bottom_up=bottom_up,
+ top_block=LastLevelP6P7_P5(out_channels, out_channels),
+ out_channels=out_channels,
+ strides=[8, 16, 32, 64, 128],
+ vit_out_dim=vit_out_dim)
+ return backbone
+
+
+@BACKBONE_REGISTRY.register()
+def build_vit_fpn_backbone_huge(cfg, input_shape: ShapeSpec):
+ window_block_indexes = (list(range(0, 7)) + list(range(8, 15)) + list(range(16, 23)) + list(range(24, 31)))
+ embed_dim = 1280
+ vit_out_dim = embed_dim
+ bottom_up = ViT( # Single-scale ViT backbone
+ img_size=1024,
+ patch_size=16,
+ embed_dim=embed_dim,
+ depth=32,
+ num_heads=16,
+ drop_path_rate=0.5,
+ window_size=14,
+ mlp_ratio=4,
+ qkv_bias=True,
+ norm_layer=partial(nn.LayerNorm, eps=1e-6),
+ window_block_indexes=window_block_indexes,
+ residual_block_indexes=[],
+ use_act_checkpoint=cfg.USE_ACT_CHECKPOINT,
+ use_rel_pos=True,
+ out_feature="last_feat",)
+
+ out_channels = cfg.MODEL.FPN.OUT_CHANNELS
+ assert out_channels == 256 or out_channels == 768 or out_channels == 1024
+ backbone = ViT_FPN(bottom_up=bottom_up,
+ top_block=LastLevelP6P7_P5(out_channels, out_channels),
+ out_channels=out_channels,
+ strides=[8, 16, 32, 64, 128],
+ vit_out_dim=vit_out_dim)
+ return backbone
diff --git a/models/grit_src/grit/modeling/meta_arch/__pycache__/grit.cpython-38.pyc b/models/grit_src/grit/modeling/meta_arch/__pycache__/grit.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..765fdea6bcc7ec42d275c07b5126d9c3bd91d129
Binary files /dev/null and b/models/grit_src/grit/modeling/meta_arch/__pycache__/grit.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/meta_arch/grit.py b/models/grit_src/grit/modeling/meta_arch/grit.py
new file mode 100644
index 0000000000000000000000000000000000000000..101725fd455e723360eaafc26db37beb226a9233
--- /dev/null
+++ b/models/grit_src/grit/modeling/meta_arch/grit.py
@@ -0,0 +1,66 @@
+from typing import Dict, List, Optional, Tuple
+import torch
+from detectron2.config import configurable
+from detectron2.structures import ImageList, Instances, Boxes
+from detectron2.modeling.meta_arch.build import META_ARCH_REGISTRY
+from detectron2.modeling.meta_arch.rcnn import GeneralizedRCNN
+
+
+@META_ARCH_REGISTRY.register()
+class GRiT(GeneralizedRCNN):
+ @configurable
+ def __init__(
+ self,
+ **kwargs):
+ super().__init__(**kwargs)
+ assert self.proposal_generator is not None
+
+ @classmethod
+ def from_config(cls, cfg):
+ ret = super().from_config(cfg)
+ return ret
+
+ def inference(
+ self,
+ batched_inputs: Tuple[Dict[str, torch.Tensor]],
+ detected_instances: Optional[List[Instances]] = None,
+ do_postprocess: bool = True,
+ ):
+ assert not self.training
+ assert detected_instances is None
+
+ images = self.preprocess_image(batched_inputs)
+ features = self.backbone(images.tensor)
+ proposals, _ = self.proposal_generator(images, features, None)
+ results, _ = self.roi_heads(features, proposals)
+ if do_postprocess:
+ assert not torch.jit.is_scripting(), \
+ "Scripting is not supported for postprocess."
+ return GRiT._postprocess(
+ results, batched_inputs, images.image_sizes)
+ else:
+ return results
+
+ def forward(self, batched_inputs: List[Dict[str, torch.Tensor]]):
+ if not self.training:
+ return self.inference(batched_inputs)
+
+ images = self.preprocess_image(batched_inputs)
+
+ gt_instances = [x["instances"].to(self.device) for x in batched_inputs]
+
+ targets_task = batched_inputs[0]['task']
+ for anno_per_image in batched_inputs:
+ assert targets_task == anno_per_image['task']
+
+ features = self.backbone(images.tensor)
+ proposals, proposal_losses = self.proposal_generator(
+ images, features, gt_instances)
+ proposals, roihead_textdecoder_losses = self.roi_heads(
+ features, proposals, gt_instances, targets_task=targets_task)
+
+ losses = {}
+ losses.update(roihead_textdecoder_losses)
+ losses.update(proposal_losses)
+
+ return losses
\ No newline at end of file
diff --git a/models/grit_src/grit/modeling/roi_heads/__pycache__/grit_fast_rcnn.cpython-38.pyc b/models/grit_src/grit/modeling/roi_heads/__pycache__/grit_fast_rcnn.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..90c9da6efaec94a1cf84d22c3ad41bb9ca69cb9d
Binary files /dev/null and b/models/grit_src/grit/modeling/roi_heads/__pycache__/grit_fast_rcnn.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/roi_heads/__pycache__/grit_roi_heads.cpython-38.pyc b/models/grit_src/grit/modeling/roi_heads/__pycache__/grit_roi_heads.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5c15041acfee7dda808343e04757a87c8b7cbc88
Binary files /dev/null and b/models/grit_src/grit/modeling/roi_heads/__pycache__/grit_roi_heads.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/roi_heads/grit_fast_rcnn.py b/models/grit_src/grit/modeling/roi_heads/grit_fast_rcnn.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d03daabac26aecf214baf1f743c97a5d7486bf7
--- /dev/null
+++ b/models/grit_src/grit/modeling/roi_heads/grit_fast_rcnn.py
@@ -0,0 +1,126 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Jialian Wu from https://github.com/facebookresearch/Detic/blob/main/detic/modeling/roi_heads/detic_fast_rcnn.py
+import torch
+from fvcore.nn import giou_loss, smooth_l1_loss
+from torch import nn
+from torch.nn import functional as F
+import fvcore.nn.weight_init as weight_init
+from detectron2.config import configurable
+from detectron2.layers import ShapeSpec, batched_nms, cat, cross_entropy, nonzero_tuple
+from detectron2.modeling.roi_heads.fast_rcnn import FastRCNNOutputLayers
+from detectron2.modeling.roi_heads.fast_rcnn import _log_classification_stats
+
+
+__all__ = ["GRiTFastRCNNOutputLayers"]
+
+
+class GRiTFastRCNNOutputLayers(FastRCNNOutputLayers):
+ @configurable
+ def __init__(
+ self,
+ input_shape: ShapeSpec,
+ **kwargs,
+ ):
+ super().__init__(
+ input_shape=input_shape,
+ **kwargs,
+ )
+
+ input_size = input_shape.channels * \
+ (input_shape.width or 1) * (input_shape.height or 1)
+
+ self.bbox_pred = nn.Sequential(
+ nn.Linear(input_size, input_size),
+ nn.ReLU(inplace=True),
+ nn.Linear(input_size, 4)
+ )
+ weight_init.c2_xavier_fill(self.bbox_pred[0])
+ nn.init.normal_(self.bbox_pred[-1].weight, std=0.001)
+ nn.init.constant_(self.bbox_pred[-1].bias, 0)
+
+ @classmethod
+ def from_config(cls, cfg, input_shape):
+ ret = super().from_config(cfg, input_shape)
+ return ret
+
+ def losses(self, predictions, proposals):
+ scores, proposal_deltas = predictions
+ gt_classes = (
+ cat([p.gt_classes for p in proposals], dim=0) if len(proposals) else torch.empty(0)
+ )
+ num_classes = self.num_classes
+ _log_classification_stats(scores, gt_classes)
+
+ if len(proposals):
+ proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0) # Nx4
+ assert not proposal_boxes.requires_grad, "Proposals should not require gradients!"
+ gt_boxes = cat(
+ [(p.gt_boxes if p.has("gt_boxes") else p.proposal_boxes).tensor for p in proposals],
+ dim=0,
+ )
+ else:
+ proposal_boxes = gt_boxes = torch.empty((0, 4), device=proposal_deltas.device)
+
+ loss_cls = self.softmax_cross_entropy_loss(scores, gt_classes)
+ return {
+ "loss_cls": loss_cls,
+ "loss_box_reg": self.box_reg_loss(
+ proposal_boxes, gt_boxes, proposal_deltas, gt_classes,
+ num_classes=num_classes)
+ }
+
+ def softmax_cross_entropy_loss(self, pred_class_logits, gt_classes):
+ if pred_class_logits.numel() == 0:
+ return pred_class_logits.new_zeros([1])[0]
+
+ loss = F.cross_entropy(
+ pred_class_logits, gt_classes, reduction="mean")
+ return loss
+
+ def box_reg_loss(
+ self, proposal_boxes, gt_boxes, pred_deltas, gt_classes,
+ num_classes=-1):
+ num_classes = num_classes if num_classes > 0 else self.num_classes
+ box_dim = proposal_boxes.shape[1]
+ fg_inds = nonzero_tuple((gt_classes >= 0) & (gt_classes < num_classes))[0]
+ if pred_deltas.shape[1] == box_dim:
+ fg_pred_deltas = pred_deltas[fg_inds]
+ else:
+ fg_pred_deltas = pred_deltas.view(-1, self.num_classes, box_dim)[
+ fg_inds, gt_classes[fg_inds]
+ ]
+
+ if self.box_reg_loss_type == "smooth_l1":
+ gt_pred_deltas = self.box2box_transform.get_deltas(
+ proposal_boxes[fg_inds],
+ gt_boxes[fg_inds],
+ )
+ loss_box_reg = smooth_l1_loss(
+ fg_pred_deltas, gt_pred_deltas, self.smooth_l1_beta, reduction="sum"
+ )
+ elif self.box_reg_loss_type == "giou":
+ fg_pred_boxes = self.box2box_transform.apply_deltas(
+ fg_pred_deltas, proposal_boxes[fg_inds]
+ )
+ loss_box_reg = giou_loss(fg_pred_boxes, gt_boxes[fg_inds], reduction="sum")
+ else:
+ raise ValueError(f"Invalid bbox reg loss type '{self.box_reg_loss_type}'")
+ return loss_box_reg / max(gt_classes.numel(), 1.0)
+
+ def predict_probs(self, predictions, proposals):
+ scores = predictions[0]
+ num_inst_per_image = [len(p) for p in proposals]
+ probs = F.softmax(scores, dim=-1)
+ return probs.split(num_inst_per_image, dim=0)
+
+ def forward(self, x):
+ if x.dim() > 2:
+ x = torch.flatten(x, start_dim=1)
+ scores = []
+
+ cls_scores = self.cls_score(x)
+ scores.append(cls_scores)
+ scores = torch.cat(scores, dim=1)
+
+ proposal_deltas = self.bbox_pred(x)
+ return scores, proposal_deltas
\ No newline at end of file
diff --git a/models/grit_src/grit/modeling/roi_heads/grit_roi_heads.py b/models/grit_src/grit/modeling/roi_heads/grit_roi_heads.py
new file mode 100644
index 0000000000000000000000000000000000000000..648214d7d3218b96e5b410350de8e52759550526
--- /dev/null
+++ b/models/grit_src/grit/modeling/roi_heads/grit_roi_heads.py
@@ -0,0 +1,478 @@
+import math
+import torch
+from typing import Dict, List, Optional, Tuple, Union
+
+from detectron2.config import configurable
+from detectron2.structures import Boxes, Instances, pairwise_iou
+from detectron2.utils.events import get_event_storage
+
+from detectron2.modeling.box_regression import Box2BoxTransform
+from detectron2.modeling.roi_heads.roi_heads import ROI_HEADS_REGISTRY, StandardROIHeads
+from detectron2.modeling.roi_heads.cascade_rcnn import CascadeROIHeads, _ScaleGradient
+from detectron2.modeling.poolers import ROIPooler
+from detectron2.layers import batched_nms
+from .grit_fast_rcnn import GRiTFastRCNNOutputLayers
+
+from ..text.text_decoder import TransformerDecoderTextualHead, GRiTTextDecoder, AutoRegressiveBeamSearch
+from ..text.load_text_token import LoadTextTokens
+from transformers import BertTokenizer
+from models.grit_src.grit.data.custom_dataset_mapper import ObjDescription
+from ..soft_nms import batched_soft_nms
+
+import logging
+logger = logging.getLogger(__name__)
+
+
+@ROI_HEADS_REGISTRY.register()
+class GRiTROIHeadsAndTextDecoder(CascadeROIHeads):
+ @configurable
+ def __init__(
+ self,
+ *,
+ text_decoder_transformer,
+ train_task: list,
+ test_task: str,
+ mult_proposal_score: bool = False,
+ mask_weight: float = 1.0,
+ object_feat_pooler=None,
+ soft_nms_enabled=False,
+ beam_size=1,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.mult_proposal_score = mult_proposal_score
+ self.mask_weight = mask_weight
+ self.object_feat_pooler = object_feat_pooler
+ self.soft_nms_enabled = soft_nms_enabled
+ self.test_task = test_task
+ self.beam_size = beam_size
+
+ tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)
+ self.tokenizer = tokenizer
+
+ assert test_task in train_task, 'GRiT has not been trained on {} task, ' \
+ 'please verify the task name or train a new ' \
+ 'GRiT on {} task'.format(test_task, test_task)
+ task_begin_tokens = {}
+ for i, task in enumerate(train_task):
+ if i == 0:
+ task_begin_tokens[task] = tokenizer.cls_token_id
+ else:
+ task_begin_tokens[task] = 103 + i
+ self.task_begin_tokens = task_begin_tokens
+
+ beamsearch_decode = AutoRegressiveBeamSearch(
+ end_token_id=tokenizer.sep_token_id,
+ max_steps=40,
+ beam_size=beam_size,
+ objectdet=test_task == "ObjectDet",
+ per_node_beam_size=1,
+ )
+ self.text_decoder = GRiTTextDecoder(
+ text_decoder_transformer,
+ beamsearch_decode=beamsearch_decode,
+ begin_token_id=task_begin_tokens[test_task],
+ loss_type='smooth',
+ tokenizer=tokenizer,
+ )
+ self.get_target_text_tokens = LoadTextTokens(tokenizer, max_text_len=40, padding='do_not_pad')
+
+ @classmethod
+ def from_config(cls, cfg, input_shape):
+ ret = super().from_config(cfg, input_shape)
+ text_decoder_transformer = TransformerDecoderTextualHead(
+ object_feature_size=cfg.MODEL.FPN.OUT_CHANNELS,
+ vocab_size=cfg.TEXT_DECODER.VOCAB_SIZE,
+ hidden_size=cfg.TEXT_DECODER.HIDDEN_SIZE,
+ num_layers=cfg.TEXT_DECODER.NUM_LAYERS,
+ attention_heads=cfg.TEXT_DECODER.ATTENTION_HEADS,
+ feedforward_size=cfg.TEXT_DECODER.FEEDFORWARD_SIZE,
+ mask_future_positions=True,
+ padding_idx=0,
+ decoder_type='bert_en',
+ use_act_checkpoint=cfg.USE_ACT_CHECKPOINT,
+ )
+ ret.update({
+ 'text_decoder_transformer': text_decoder_transformer,
+ 'train_task': cfg.MODEL.TRAIN_TASK,
+ 'test_task': cfg.MODEL.TEST_TASK,
+ 'mult_proposal_score': cfg.MODEL.ROI_BOX_HEAD.MULT_PROPOSAL_SCORE,
+ 'mask_weight': cfg.MODEL.ROI_HEADS.MASK_WEIGHT,
+ 'soft_nms_enabled': cfg.MODEL.ROI_HEADS.SOFT_NMS_ENABLED,
+ 'beam_size': cfg.MODEL.BEAM_SIZE,
+ })
+ return ret
+
+ @classmethod
+ def _init_box_head(self, cfg, input_shape):
+ ret = super()._init_box_head(cfg, input_shape)
+ del ret['box_predictors']
+ cascade_bbox_reg_weights = cfg.MODEL.ROI_BOX_CASCADE_HEAD.BBOX_REG_WEIGHTS
+ box_predictors = []
+ for box_head, bbox_reg_weights in zip(ret['box_heads'], \
+ cascade_bbox_reg_weights):
+ box_predictors.append(
+ GRiTFastRCNNOutputLayers(
+ cfg, box_head.output_shape,
+ box2box_transform=Box2BoxTransform(weights=bbox_reg_weights)
+ ))
+ ret['box_predictors'] = box_predictors
+
+ in_features = cfg.MODEL.ROI_HEADS.IN_FEATURES
+ pooler_scales = tuple(1.0 / input_shape[k].stride for k in in_features)
+ sampling_ratio = cfg.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
+ pooler_type = cfg.MODEL.ROI_BOX_HEAD.POOLER_TYPE
+ object_feat_pooler = ROIPooler(
+ output_size=cfg.MODEL.ROI_HEADS.OBJECT_FEAT_POOLER_RES,
+ scales=pooler_scales,
+ sampling_ratio=sampling_ratio,
+ pooler_type=pooler_type,
+ )
+ ret['object_feat_pooler'] = object_feat_pooler
+ return ret
+
+ def check_if_all_background(self, proposals, targets, stage):
+ all_background = True
+ for proposals_per_image in proposals:
+ if not (proposals_per_image.gt_classes == self.num_classes).all():
+ all_background = False
+
+ if all_background:
+ logger.info('all proposals are background at stage {}'.format(stage))
+ proposals[0].proposal_boxes.tensor[0, :] = targets[0].gt_boxes.tensor[0, :]
+ proposals[0].gt_boxes.tensor[0, :] = targets[0].gt_boxes.tensor[0, :]
+ proposals[0].objectness_logits[0] = math.log((1.0 - 1e-10) / (1 - (1.0 - 1e-10)))
+ proposals[0].gt_classes[0] = targets[0].gt_classes[0]
+ proposals[0].gt_object_descriptions.data[0] = targets[0].gt_object_descriptions.data[0]
+ if 'foreground' in proposals[0].get_fields().keys():
+ proposals[0].foreground[0] = 1
+ return proposals
+
+ def _forward_box(self, features, proposals, targets=None, task="ObjectDet"):
+ if self.training:
+ proposals = self.check_if_all_background(proposals, targets, 0)
+ if (not self.training) and self.mult_proposal_score:
+ if len(proposals) > 0 and proposals[0].has('scores'):
+ proposal_scores = [p.get('scores') for p in proposals]
+ else:
+ proposal_scores = [p.get('objectness_logits') for p in proposals]
+
+ features = [features[f] for f in self.box_in_features]
+ head_outputs = []
+ prev_pred_boxes = None
+ image_sizes = [x.image_size for x in proposals]
+
+ for k in range(self.num_cascade_stages):
+ if k > 0:
+ proposals = self._create_proposals_from_boxes(
+ prev_pred_boxes, image_sizes,
+ logits=[p.objectness_logits for p in proposals])
+ if self.training:
+ proposals = self._match_and_label_boxes_GRiT(
+ proposals, k, targets)
+ proposals = self.check_if_all_background(proposals, targets, k)
+ predictions = self._run_stage(features, proposals, k)
+ prev_pred_boxes = self.box_predictor[k].predict_boxes(
+ (predictions[0], predictions[1]), proposals)
+ head_outputs.append((self.box_predictor[k], predictions, proposals))
+
+ if self.training:
+ object_features = self.object_feat_pooler(features, [x.proposal_boxes for x in proposals])
+ object_features = _ScaleGradient.apply(object_features, 1.0 / self.num_cascade_stages)
+ foreground = torch.cat([x.foreground for x in proposals])
+ object_features = object_features[foreground > 0]
+
+ object_descriptions = []
+ for x in proposals:
+ object_descriptions += x.gt_object_descriptions[x.foreground > 0].data
+ object_descriptions = ObjDescription(object_descriptions)
+ object_descriptions = object_descriptions.data
+
+ if len(object_descriptions) > 0:
+ begin_token = self.task_begin_tokens[task]
+ text_decoder_inputs = self.get_target_text_tokens(object_descriptions, object_features, begin_token)
+ object_features = object_features.view(
+ object_features.shape[0], object_features.shape[1], -1).permute(0, 2, 1).contiguous()
+ text_decoder_inputs.update({'object_features': object_features})
+ text_decoder_loss = self.text_decoder(text_decoder_inputs)
+ else:
+ text_decoder_loss = head_outputs[0][1][0].new_zeros([1])[0]
+
+ losses = {}
+ storage = get_event_storage()
+ # RoI Head losses (For the proposal generator loss, please find it in grit.py)
+ for stage, (predictor, predictions, proposals) in enumerate(head_outputs):
+ with storage.name_scope("stage{}".format(stage)):
+ stage_losses = predictor.losses(
+ (predictions[0], predictions[1]), proposals)
+ losses.update({k + "_stage{}".format(stage): v for k, v in stage_losses.items()})
+ # Text Decoder loss
+ losses.update({'text_decoder_loss': text_decoder_loss})
+ return losses
+ else:
+ scores_per_stage = [h[0].predict_probs(h[1], h[2]) for h in head_outputs]
+ logits_per_stage = [(h[1][0],) for h in head_outputs]
+ scores = [
+ sum(list(scores_per_image)) * (1.0 / self.num_cascade_stages)
+ for scores_per_image in zip(*scores_per_stage)
+ ]
+ logits = [
+ sum(list(logits_per_image)) * (1.0 / self.num_cascade_stages)
+ for logits_per_image in zip(*logits_per_stage)
+ ]
+ if self.mult_proposal_score:
+ scores = [(s * ps[:, None]) ** 0.5 for s, ps in zip(scores, proposal_scores)]
+ predictor, predictions, proposals = head_outputs[-1]
+ boxes = predictor.predict_boxes(
+ (predictions[0], predictions[1]), proposals)
+ assert len(boxes) == 1
+ pred_instances, _ = self.fast_rcnn_inference_GRiT(
+ boxes,
+ scores,
+ logits,
+ image_sizes,
+ predictor.test_score_thresh,
+ predictor.test_nms_thresh,
+ predictor.test_topk_per_image,
+ self.soft_nms_enabled,
+ )
+
+ assert len(pred_instances) == 1, "Only support one image"
+ for i, pred_instance in enumerate(pred_instances):
+ if len(pred_instance.pred_boxes) > 0:
+ object_features = self.object_feat_pooler(features, [pred_instance.pred_boxes])
+ object_features = object_features.view(
+ object_features.shape[0], object_features.shape[1], -1).permute(0, 2, 1).contiguous()
+ text_decoder_output = self.text_decoder({'object_features': object_features})
+ if self.beam_size > 1 and self.test_task == "ObjectDet":
+ pred_boxes = []
+ pred_scores = []
+ pred_classes = []
+ pred_object_descriptions = []
+
+ for beam_id in range(self.beam_size):
+ pred_boxes.append(pred_instance.pred_boxes.tensor)
+ # object score = sqrt(objectness score x description score)
+ pred_scores.append((pred_instance.scores *
+ torch.exp(text_decoder_output['logprobs'])[:, beam_id]) ** 0.5)
+ pred_classes.append(pred_instance.pred_classes)
+ for prediction in text_decoder_output['predictions'][:, beam_id, :]:
+ # convert text tokens to words
+ description = self.tokenizer.decode(prediction.tolist()[1:], skip_special_tokens=True)
+ pred_object_descriptions.append(description)
+
+ merged_instances = Instances(image_sizes[0])
+ if torch.cat(pred_scores, dim=0).shape[0] <= predictor.test_topk_per_image:
+ merged_instances.scores = torch.cat(pred_scores, dim=0)
+ merged_instances.pred_boxes = Boxes(torch.cat(pred_boxes, dim=0))
+ merged_instances.pred_classes = torch.cat(pred_classes, dim=0)
+ merged_instances.pred_object_descriptions = ObjDescription(pred_object_descriptions)
+ else:
+ pred_scores, top_idx = torch.topk(
+ torch.cat(pred_scores, dim=0), predictor.test_topk_per_image)
+ merged_instances.scores = pred_scores
+ merged_instances.pred_boxes = Boxes(torch.cat(pred_boxes, dim=0)[top_idx, :])
+ merged_instances.pred_classes = torch.cat(pred_classes, dim=0)[top_idx]
+ merged_instances.pred_object_descriptions = \
+ ObjDescription(ObjDescription(pred_object_descriptions)[top_idx].data)
+
+ pred_instances[i] = merged_instances
+ else:
+ # object score = sqrt(objectness score x description score)
+ pred_instance.scores = (pred_instance.scores *
+ torch.exp(text_decoder_output['logprobs'])) ** 0.5
+
+ pred_object_descriptions = []
+ for prediction in text_decoder_output['predictions']:
+ # convert text tokens to words
+ description = self.tokenizer.decode(prediction.tolist()[1:], skip_special_tokens=True)
+ pred_object_descriptions.append(description)
+ pred_instance.pred_object_descriptions = ObjDescription(pred_object_descriptions)
+ else:
+ pred_instance.pred_object_descriptions = ObjDescription([])
+
+ return pred_instances
+
+
+ def forward(self, features, proposals, targets=None, targets_task="ObjectDet"):
+ if self.training:
+ proposals = self.label_and_sample_proposals(
+ proposals, targets)
+
+ losses = self._forward_box(features, proposals, targets, task=targets_task)
+ if targets[0].has('gt_masks'):
+ mask_losses = self._forward_mask(features, proposals)
+ losses.update({k: v * self.mask_weight \
+ for k, v in mask_losses.items()})
+ else:
+ losses.update(self._get_empty_mask_loss(device=proposals[0].objectness_logits.device))
+ return proposals, losses
+ else:
+ pred_instances = self._forward_box(features, proposals, task=self.test_task)
+ pred_instances = self.forward_with_given_boxes(features, pred_instances)
+ return pred_instances, {}
+
+ @torch.no_grad()
+ def _match_and_label_boxes_GRiT(self, proposals, stage, targets):
+ """
+ Add "gt_object_description" and "foreground" to detectron2's _match_and_label_boxes
+ """
+ num_fg_samples, num_bg_samples = [], []
+ for proposals_per_image, targets_per_image in zip(proposals, targets):
+ match_quality_matrix = pairwise_iou(
+ targets_per_image.gt_boxes, proposals_per_image.proposal_boxes
+ )
+ # proposal_labels are 0 or 1
+ matched_idxs, proposal_labels = self.proposal_matchers[stage](match_quality_matrix)
+ if len(targets_per_image) > 0:
+ gt_classes = targets_per_image.gt_classes[matched_idxs]
+ # Label unmatched proposals (0 label from matcher) as background (label=num_classes)
+ gt_classes[proposal_labels == 0] = self.num_classes
+ foreground = torch.ones_like(gt_classes)
+ foreground[proposal_labels == 0] = 0
+ gt_boxes = targets_per_image.gt_boxes[matched_idxs]
+ gt_object_descriptions = targets_per_image.gt_object_descriptions[matched_idxs]
+ else:
+ gt_classes = torch.zeros_like(matched_idxs) + self.num_classes
+ foreground = torch.zeros_like(gt_classes)
+ gt_boxes = Boxes(
+ targets_per_image.gt_boxes.tensor.new_zeros((len(proposals_per_image), 4))
+ )
+ gt_object_descriptions = ObjDescription(['None' for i in range(len(proposals_per_image))])
+ proposals_per_image.gt_classes = gt_classes
+ proposals_per_image.gt_boxes = gt_boxes
+ proposals_per_image.gt_object_descriptions = gt_object_descriptions
+ proposals_per_image.foreground = foreground
+
+ num_fg_samples.append((proposal_labels == 1).sum().item())
+ num_bg_samples.append(proposal_labels.numel() - num_fg_samples[-1])
+
+ # Log the number of fg/bg samples in each stage
+ storage = get_event_storage()
+ storage.put_scalar(
+ "stage{}/roi_head/num_fg_samples".format(stage),
+ sum(num_fg_samples) / len(num_fg_samples),
+ )
+ storage.put_scalar(
+ "stage{}/roi_head/num_bg_samples".format(stage),
+ sum(num_bg_samples) / len(num_bg_samples),
+ )
+ return proposals
+
+ def fast_rcnn_inference_GRiT(
+ self,
+ boxes: List[torch.Tensor],
+ scores: List[torch.Tensor],
+ logits: List[torch.Tensor],
+ image_shapes: List[Tuple[int, int]],
+ score_thresh: float,
+ nms_thresh: float,
+ topk_per_image: int,
+ soft_nms_enabled: bool,
+ ):
+ result_per_image = [
+ self.fast_rcnn_inference_single_image_GRiT(
+ boxes_per_image, scores_per_image, logits_per_image, image_shape,
+ score_thresh, nms_thresh, topk_per_image, soft_nms_enabled
+ )
+ for scores_per_image, boxes_per_image, image_shape, logits_per_image \
+ in zip(scores, boxes, image_shapes, logits)
+ ]
+ return [x[0] for x in result_per_image], [x[1] for x in result_per_image]
+
+ def fast_rcnn_inference_single_image_GRiT(
+ self,
+ boxes,
+ scores,
+ logits,
+ image_shape: Tuple[int, int],
+ score_thresh: float,
+ nms_thresh: float,
+ topk_per_image: int,
+ soft_nms_enabled,
+ ):
+ """
+ Add soft NMS to detectron2's fast_rcnn_inference_single_image
+ """
+ valid_mask = torch.isfinite(boxes).all(dim=1) & torch.isfinite(scores).all(dim=1)
+ if not valid_mask.all():
+ boxes = boxes[valid_mask]
+ scores = scores[valid_mask]
+ logits = logits[valid_mask]
+
+ scores = scores[:, :-1]
+ logits = logits[:, :-1]
+ num_bbox_reg_classes = boxes.shape[1] // 4
+ # Convert to Boxes to use the `clip` function ...
+ boxes = Boxes(boxes.reshape(-1, 4))
+ boxes.clip(image_shape)
+ boxes = boxes.tensor.view(-1, num_bbox_reg_classes, 4) # R x C x 4
+
+ # 1. Filter results based on detection scores. It can make NMS more efficient
+ # by filtering out low-confidence detections.
+ filter_mask = scores > score_thresh # R x K
+ # R' x 2. First column contains indices of the R predictions;
+ # Second column contains indices of classes.
+ filter_inds = filter_mask.nonzero()
+ if num_bbox_reg_classes == 1:
+ boxes = boxes[filter_inds[:, 0], 0]
+ else:
+ boxes = boxes[filter_mask]
+ scores = scores[filter_mask]
+ logits = logits[filter_mask]
+
+ # 2. Apply NMS for each class independently.
+ if not soft_nms_enabled:
+ keep = batched_nms(boxes, scores, filter_inds[:, 1], nms_thresh)
+ else:
+ keep, soft_nms_scores = batched_soft_nms(
+ boxes,
+ scores,
+ filter_inds[:, 1],
+ "linear",
+ 0.5,
+ nms_thresh,
+ 0.001,
+ )
+ scores[keep] = soft_nms_scores
+ if topk_per_image >= 0:
+ keep = keep[:topk_per_image]
+ boxes, scores, filter_inds = boxes[keep], scores[keep], filter_inds[keep]
+ logits = logits[keep]
+
+ result = Instances(image_shape)
+ result.pred_boxes = Boxes(boxes)
+ result.scores = scores
+ result.pred_classes = filter_inds[:, 1]
+ result.logits = logits
+ return result, filter_inds[:, 0]
+
+ def _get_empty_mask_loss(self, device):
+ if self.mask_on:
+ return {'loss_mask': torch.zeros(
+ (1, ), device=device, dtype=torch.float32)[0]}
+ else:
+ return {}
+
+ def _create_proposals_from_boxes(self, boxes, image_sizes, logits):
+ boxes = [Boxes(b.detach()) for b in boxes]
+ proposals = []
+ for boxes_per_image, image_size, logit in zip(
+ boxes, image_sizes, logits):
+ boxes_per_image.clip(image_size)
+ if self.training:
+ inds = boxes_per_image.nonempty()
+ boxes_per_image = boxes_per_image[inds]
+ logit = logit[inds]
+ prop = Instances(image_size)
+ prop.proposal_boxes = boxes_per_image
+ prop.objectness_logits = logit
+ proposals.append(prop)
+ return proposals
+
+ def _run_stage(self, features, proposals, stage):
+ pool_boxes = [x.proposal_boxes for x in proposals]
+ box_features = self.box_pooler(features, pool_boxes)
+ box_features = _ScaleGradient.apply(box_features, 1.0 / self.num_cascade_stages)
+ box_features = self.box_head[stage](box_features)
+ return self.box_predictor[stage](box_features)
diff --git a/models/grit_src/grit/modeling/soft_nms.py b/models/grit_src/grit/modeling/soft_nms.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a5aae7c4261191b8e07e0fd25055d8917f7f97d
--- /dev/null
+++ b/models/grit_src/grit/modeling/soft_nms.py
@@ -0,0 +1,177 @@
+import torch
+
+from detectron2.structures import Boxes, RotatedBoxes, pairwise_iou, pairwise_iou_rotated
+
+
+def soft_nms(boxes, scores, method, gaussian_sigma, linear_threshold, prune_threshold):
+ """
+ Performs soft non-maximum suppression algorithm on axis aligned boxes
+
+ Args:
+ boxes (Tensor[N, 5]):
+ boxes where NMS will be performed. They
+ are expected to be in (x_ctr, y_ctr, width, height, angle_degrees) format
+ scores (Tensor[N]):
+ scores for each one of the boxes
+ method (str):
+ one of ['gaussian', 'linear', 'hard']
+ see paper for details. users encouraged not to use "hard", as this is the
+ same nms available elsewhere in detectron2
+ gaussian_sigma (float):
+ parameter for Gaussian penalty function
+ linear_threshold (float):
+ iou threshold for applying linear decay. Nt from the paper
+ re-used as threshold for standard "hard" nms
+ prune_threshold (float):
+ boxes with scores below this threshold are pruned at each iteration.
+ Dramatically reduces computation time. Authors use values in [10e-4, 10e-2]
+
+ Returns:
+ tuple(Tensor, Tensor):
+ [0]: int64 tensor with the indices of the elements that have been kept
+ by Soft NMS, sorted in decreasing order of scores
+ [1]: float tensor with the re-scored scores of the elements that were kept
+"""
+ return _soft_nms(
+ Boxes,
+ pairwise_iou,
+ boxes,
+ scores,
+ method,
+ gaussian_sigma,
+ linear_threshold,
+ prune_threshold,
+ )
+
+
+def batched_soft_nms(
+ boxes, scores, idxs, method, gaussian_sigma, linear_threshold, prune_threshold
+):
+ """
+ Performs soft non-maximum suppression in a batched fashion.
+
+ Each index value correspond to a category, and NMS
+ will not be applied between elements of different categories.
+
+ Args:
+ boxes (Tensor[N, 4]):
+ boxes where NMS will be performed. They
+ are expected to be in (x1, y1, x2, y2) format
+ scores (Tensor[N]):
+ scores for each one of the boxes
+ idxs (Tensor[N]):
+ indices of the categories for each one of the boxes.
+ method (str):
+ one of ['gaussian', 'linear', 'hard']
+ see paper for details. users encouraged not to use "hard", as this is the
+ same nms available elsewhere in detectron2
+ gaussian_sigma (float):
+ parameter for Gaussian penalty function
+ linear_threshold (float):
+ iou threshold for applying linear decay. Nt from the paper
+ re-used as threshold for standard "hard" nms
+ prune_threshold (float):
+ boxes with scores below this threshold are pruned at each iteration.
+ Dramatically reduces computation time. Authors use values in [10e-4, 10e-2]
+ Returns:
+ tuple(Tensor, Tensor):
+ [0]: int64 tensor with the indices of the elements that have been kept
+ by Soft NMS, sorted in decreasing order of scores
+ [1]: float tensor with the re-scored scores of the elements that were kept
+ """
+ if boxes.numel() == 0:
+ return (
+ torch.empty((0,), dtype=torch.int64, device=boxes.device),
+ torch.empty((0,), dtype=torch.float32, device=scores.device),
+ )
+ # strategy: in order to perform NMS independently per class.
+ # we add an offset to all the boxes. The offset is dependent
+ # only on the class idx, and is large enough so that boxes
+ # from different classes do not overlap
+ max_coordinate = boxes.max()
+ offsets = idxs.to(boxes) * (max_coordinate + 1)
+ boxes_for_nms = boxes + offsets[:, None]
+ return soft_nms(
+ boxes_for_nms, scores, method, gaussian_sigma, linear_threshold, prune_threshold
+ )
+
+
+def _soft_nms(
+ box_class,
+ pairwise_iou_func,
+ boxes,
+ scores,
+ method,
+ gaussian_sigma,
+ linear_threshold,
+ prune_threshold,
+):
+ """
+ Soft non-max suppression algorithm.
+
+ Implementation of [Soft-NMS -- Improving Object Detection With One Line of Codec]
+ (https://arxiv.org/abs/1704.04503)
+
+ Args:
+ box_class (cls): one of Box, RotatedBoxes
+ pairwise_iou_func (func): one of pairwise_iou, pairwise_iou_rotated
+ boxes (Tensor[N, ?]):
+ boxes where NMS will be performed
+ if Boxes, in (x1, y1, x2, y2) format
+ if RotatedBoxes, in (x_ctr, y_ctr, width, height, angle_degrees) format
+ scores (Tensor[N]):
+ scores for each one of the boxes
+ method (str):
+ one of ['gaussian', 'linear', 'hard']
+ see paper for details. users encouraged not to use "hard", as this is the
+ same nms available elsewhere in detectron2
+ gaussian_sigma (float):
+ parameter for Gaussian penalty function
+ linear_threshold (float):
+ iou threshold for applying linear decay. Nt from the paper
+ re-used as threshold for standard "hard" nms
+ prune_threshold (float):
+ boxes with scores below this threshold are pruned at each iteration.
+ Dramatically reduces computation time. Authors use values in [10e-4, 10e-2]
+
+ Returns:
+ tuple(Tensor, Tensor):
+ [0]: int64 tensor with the indices of the elements that have been kept
+ by Soft NMS, sorted in decreasing order of scores
+ [1]: float tensor with the re-scored scores of the elements that were kept
+ """
+ boxes = boxes.clone()
+ scores = scores.clone()
+ idxs = torch.arange(scores.size()[0])
+
+ idxs_out = []
+ scores_out = []
+
+ while scores.numel() > 0:
+ top_idx = torch.argmax(scores)
+ idxs_out.append(idxs[top_idx].item())
+ scores_out.append(scores[top_idx].item())
+
+ top_box = boxes[top_idx]
+ ious = pairwise_iou_func(box_class(top_box.unsqueeze(0)), box_class(boxes))[0]
+
+ if method == "linear":
+ decay = torch.ones_like(ious)
+ decay_mask = ious > linear_threshold
+ decay[decay_mask] = 1 - ious[decay_mask]
+ elif method == "gaussian":
+ decay = torch.exp(-torch.pow(ious, 2) / gaussian_sigma)
+ elif method == "hard": # standard NMS
+ decay = (ious < linear_threshold).float()
+ else:
+ raise NotImplementedError("{} soft nms method not implemented.".format(method))
+
+ scores *= decay
+ keep = scores > prune_threshold
+ keep[top_idx] = False
+
+ boxes = boxes[keep]
+ scores = scores[keep]
+ idxs = idxs[keep]
+
+ return torch.tensor(idxs_out).to(boxes.device), torch.tensor(scores_out).to(scores.device)
\ No newline at end of file
diff --git a/models/grit_src/grit/modeling/text/__pycache__/file_utils.cpython-38.pyc b/models/grit_src/grit/modeling/text/__pycache__/file_utils.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7aaee6471a10beb3736b3901a6915c3be584881c
Binary files /dev/null and b/models/grit_src/grit/modeling/text/__pycache__/file_utils.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/text/__pycache__/load_text_token.cpython-38.pyc b/models/grit_src/grit/modeling/text/__pycache__/load_text_token.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..46957825e3a97d98628d000cc90da0454ec3530b
Binary files /dev/null and b/models/grit_src/grit/modeling/text/__pycache__/load_text_token.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/text/__pycache__/modeling_bert.cpython-38.pyc b/models/grit_src/grit/modeling/text/__pycache__/modeling_bert.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1769fa39d6b8f84b2e6b4fe1b5e658e02e55f4de
Binary files /dev/null and b/models/grit_src/grit/modeling/text/__pycache__/modeling_bert.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/text/__pycache__/text_decoder.cpython-38.pyc b/models/grit_src/grit/modeling/text/__pycache__/text_decoder.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..487bfd4fed5d94e9cb5d967721f69a3dd8b1b9a9
Binary files /dev/null and b/models/grit_src/grit/modeling/text/__pycache__/text_decoder.cpython-38.pyc differ
diff --git a/models/grit_src/grit/modeling/text/file_utils.py b/models/grit_src/grit/modeling/text/file_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..51918cf3857471e4ffb5b617d73ee8b9eed0989e
--- /dev/null
+++ b/models/grit_src/grit/modeling/text/file_utils.py
@@ -0,0 +1,256 @@
+# Utilities for working with the local dataset cache.
+# This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp
+# Copyright by the AllenNLP authors.
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import sys
+import json
+import logging
+import os
+import shutil
+import tempfile
+import fnmatch
+from functools import wraps
+from hashlib import sha256
+from io import open
+
+import boto3
+import requests
+from botocore.exceptions import ClientError
+from tqdm import tqdm
+
+try:
+ from torch.hub import _get_torch_home
+ torch_cache_home = _get_torch_home()
+except ImportError:
+ torch_cache_home = os.path.expanduser(
+ os.getenv('TORCH_HOME', os.path.join(
+ os.getenv('XDG_CACHE_HOME', '~/.cache'), 'torch')))
+default_cache_path = os.path.join(torch_cache_home, 'pytorch_transformers')
+
+try:
+ from urllib.parse import urlparse
+except ImportError:
+ from urlparse import urlparse
+
+try:
+ from pathlib import Path
+ PYTORCH_PRETRAINED_BERT_CACHE = Path(
+ os.getenv('PYTORCH_PRETRAINED_BERT_CACHE', default_cache_path))
+except (AttributeError, ImportError):
+ PYTORCH_PRETRAINED_BERT_CACHE = os.getenv('PYTORCH_PRETRAINED_BERT_CACHE',
+ default_cache_path)
+
+logger = logging.getLogger(__name__) # pylint: disable=invalid-name
+
+
+def url_to_filename(url, etag=None):
+ """
+ Convert `url` into a hashed filename in a repeatable way.
+ If `etag` is specified, append its hash to the url's, delimited
+ by a period.
+ """
+ url_bytes = url.encode('utf-8')
+ url_hash = sha256(url_bytes)
+ filename = url_hash.hexdigest()
+
+ if etag:
+ etag_bytes = etag.encode('utf-8')
+ etag_hash = sha256(etag_bytes)
+ filename += '.' + etag_hash.hexdigest()
+
+ return filename
+
+
+def filename_to_url(filename, cache_dir=None):
+ """
+ Return the url and etag (which may be ``None``) stored for `filename`.
+ Raise ``EnvironmentError`` if `filename` or its stored metadata do not exist.
+ """
+ if cache_dir is None:
+ cache_dir = PYTORCH_PRETRAINED_BERT_CACHE
+ if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+ cache_dir = str(cache_dir)
+
+ cache_path = os.path.join(cache_dir, filename)
+ if not os.path.exists(cache_path):
+ raise EnvironmentError("file {} not found".format(cache_path))
+
+ meta_path = cache_path + '.json'
+ if not os.path.exists(meta_path):
+ raise EnvironmentError("file {} not found".format(meta_path))
+
+ with open(meta_path, encoding="utf-8") as meta_file:
+ metadata = json.load(meta_file)
+ url = metadata['url']
+ etag = metadata['etag']
+
+ return url, etag
+
+
+def cached_path(url_or_filename, cache_dir=None):
+ """
+ Given something that might be a URL (or might be a local path),
+ determine which. If it's a URL, download the file and cache it, and
+ return the path to the cached file. If it's already a local path,
+ make sure the file exists and then return the path.
+ """
+ if cache_dir is None:
+ cache_dir = PYTORCH_PRETRAINED_BERT_CACHE
+ if sys.version_info[0] == 3 and isinstance(url_or_filename, Path):
+ url_or_filename = str(url_or_filename)
+ if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+ cache_dir = str(cache_dir)
+
+ parsed = urlparse(url_or_filename)
+
+ if parsed.scheme in ('http', 'https', 's3'):
+ # URL, so get it from the cache (downloading if necessary)
+ return get_from_cache(url_or_filename, cache_dir)
+ elif os.path.exists(url_or_filename):
+ # File, and it exists.
+ return url_or_filename
+ elif parsed.scheme == '':
+ # File, but it doesn't exist.
+ raise EnvironmentError("file {} not found".format(url_or_filename))
+ else:
+ # Something unknown
+ raise ValueError("unable to parse {} as a URL or as a local path".format(url_or_filename))
+
+
+def split_s3_path(url):
+ """Split a full s3 path into the bucket name and path."""
+ parsed = urlparse(url)
+ if not parsed.netloc or not parsed.path:
+ raise ValueError("bad s3 path {}".format(url))
+ bucket_name = parsed.netloc
+ s3_path = parsed.path
+ # Remove '/' at beginning of path.
+ if s3_path.startswith("/"):
+ s3_path = s3_path[1:]
+ return bucket_name, s3_path
+
+
+def s3_request(func):
+ """
+ Wrapper function for s3 requests in order to create more helpful error
+ messages.
+ """
+
+ @wraps(func)
+ def wrapper(url, *args, **kwargs):
+ try:
+ return func(url, *args, **kwargs)
+ except ClientError as exc:
+ if int(exc.response["Error"]["Code"]) == 404:
+ raise EnvironmentError("file {} not found".format(url))
+ else:
+ raise
+
+ return wrapper
+
+
+@s3_request
+def s3_etag(url):
+ """Check ETag on S3 object."""
+ s3_resource = boto3.resource("s3")
+ bucket_name, s3_path = split_s3_path(url)
+ s3_object = s3_resource.Object(bucket_name, s3_path)
+ return s3_object.e_tag
+
+
+@s3_request
+def s3_get(url, temp_file):
+ """Pull a file directly from S3."""
+ s3_resource = boto3.resource("s3")
+ bucket_name, s3_path = split_s3_path(url)
+ s3_resource.Bucket(bucket_name).download_fileobj(s3_path, temp_file)
+
+
+def http_get(url, temp_file):
+ req = requests.get(url, stream=True)
+ content_length = req.headers.get('Content-Length')
+ total = int(content_length) if content_length is not None else None
+ progress = tqdm(unit="B", total=total)
+ for chunk in req.iter_content(chunk_size=1024):
+ if chunk: # filter out keep-alive new chunks
+ progress.update(len(chunk))
+ temp_file.write(chunk)
+ progress.close()
+
+
+def get_from_cache(url, cache_dir=None):
+ """
+ Given a URL, look for the corresponding dataset in the local cache.
+ If it's not there, download it. Then return the path to the cached file.
+ """
+ if cache_dir is None:
+ cache_dir = PYTORCH_PRETRAINED_BERT_CACHE
+ if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+ cache_dir = str(cache_dir)
+ if sys.version_info[0] == 2 and not isinstance(cache_dir, str):
+ cache_dir = str(cache_dir)
+
+ if not os.path.exists(cache_dir):
+ os.makedirs(cache_dir)
+
+ # Get eTag to add to filename, if it exists.
+ if url.startswith("s3://"):
+ etag = s3_etag(url)
+ else:
+ try:
+ response = requests.head(url, allow_redirects=True)
+ if response.status_code != 200:
+ etag = None
+ else:
+ etag = response.headers.get("ETag")
+ except EnvironmentError:
+ etag = None
+
+ if sys.version_info[0] == 2 and etag is not None:
+ etag = etag.decode('utf-8')
+ filename = url_to_filename(url, etag)
+
+ # get cache path to put the file
+ cache_path = os.path.join(cache_dir, filename)
+
+ # If we don't have a connection (etag is None) and can't identify the file
+ # try to get the last downloaded one
+ if not os.path.exists(cache_path) and etag is None:
+ matching_files = fnmatch.filter(os.listdir(cache_dir), filename + '.*')
+ matching_files = list(filter(lambda s: not s.endswith('.json'), matching_files))
+ if matching_files:
+ cache_path = os.path.join(cache_dir, matching_files[-1])
+
+ if not os.path.exists(cache_path):
+ # Download to temporary file, then copy to cache dir once finished.
+ # Otherwise you get corrupt cache entries if the download gets interrupted.
+ with tempfile.NamedTemporaryFile() as temp_file:
+ logger.info("%s not found in cache, downloading to %s", url, temp_file.name)
+
+ # GET file object
+ if url.startswith("s3://"):
+ s3_get(url, temp_file)
+ else:
+ http_get(url, temp_file)
+
+ # we are copying the file before closing it, so flush to avoid truncation
+ temp_file.flush()
+ # shutil.copyfileobj() starts at the current position, so go to the start
+ temp_file.seek(0)
+
+ logger.info("copying %s to cache at %s", temp_file.name, cache_path)
+ with open(cache_path, 'wb') as cache_file:
+ shutil.copyfileobj(temp_file, cache_file)
+
+ logger.info("creating metadata file for %s", cache_path)
+ meta = {'url': url, 'etag': etag}
+ meta_path = cache_path + '.json'
+ with open(meta_path, 'w') as meta_file:
+ output_string = json.dumps(meta)
+ meta_file.write(output_string)
+
+ logger.info("removing temp file %s", temp_file.name)
+
+ return cache_path
diff --git a/models/grit_src/grit/modeling/text/load_text_token.py b/models/grit_src/grit/modeling/text/load_text_token.py
new file mode 100644
index 0000000000000000000000000000000000000000..8491021bf5d7d23d7f3826395f270dccad30df36
--- /dev/null
+++ b/models/grit_src/grit/modeling/text/load_text_token.py
@@ -0,0 +1,80 @@
+import torch
+
+
+class LoadTextTokens(object):
+ def __init__(self, tokenizer, max_text_len=40, padding='do_not_pad'):
+ self.tokenizer = tokenizer
+ self.max_text_len = max_text_len
+ self.padding = padding
+
+ def descriptions_to_text_tokens(self, target, begin_token):
+ target_encoding = self.tokenizer(
+ target, padding=self.padding,
+ add_special_tokens=False,
+ truncation=True, max_length=self.max_text_len)
+
+ need_predict = [1] * len(target_encoding['input_ids'])
+ payload = target_encoding['input_ids']
+ if len(payload) > self.max_text_len - 2:
+ payload = payload[-(self.max_text_len - 2):]
+ need_predict = payload[-(self.max_text_len - 2):]
+
+ input_ids = [begin_token] + payload + [self.tokenizer.sep_token_id]
+
+ need_predict = [0] + need_predict + [1]
+ data = {
+ 'text_tokens': torch.tensor(input_ids),
+ 'text_lengths': len(input_ids),
+ 'need_predict': torch.tensor(need_predict),
+ }
+
+ return data
+
+ def __call__(self, object_descriptions, box_features, begin_token):
+ text_tokens = []
+ text_lengths = []
+ need_predict = []
+ for description in object_descriptions:
+ tokens = self.descriptions_to_text_tokens(description, begin_token)
+ text_tokens.append(tokens['text_tokens'])
+ text_lengths.append(tokens['text_lengths'])
+ need_predict.append(tokens['need_predict'])
+
+ text_tokens = torch.cat(self.collate(text_tokens), dim=0).to(box_features.device)
+ text_lengths = torch.tensor(text_lengths).to(box_features.device)
+ need_predict = torch.cat(self.collate(need_predict), dim=0).to(box_features.device)
+
+ assert text_tokens.dim() == 2 and need_predict.dim() == 2
+ data = {'text_tokens': text_tokens,
+ 'text_lengths': text_lengths,
+ 'need_predict': need_predict}
+
+ return data
+
+ def collate(self, batch):
+ if all(isinstance(b, torch.Tensor) for b in batch) and len(batch) > 0:
+ if not all(b.shape == batch[0].shape for b in batch[1:]):
+ assert all(len(b.shape) == len(batch[0].shape) for b in batch[1:])
+ shape = torch.tensor([b.shape for b in batch])
+ max_shape = tuple(shape.max(dim=0)[0].tolist())
+ batch2 = []
+ for b in batch:
+ if any(c < m for c, m in zip(b.shape, max_shape)):
+ b2 = torch.zeros(max_shape, dtype=b.dtype, device=b.device)
+ if b.dim() == 1:
+ b2[:b.shape[0]] = b
+ elif b.dim() == 2:
+ b2[:b.shape[0], :b.shape[1]] = b
+ elif b.dim() == 3:
+ b2[:b.shape[0], :b.shape[1], :b.shape[2]] = b
+ else:
+ raise NotImplementedError
+ b = b2
+ batch2.append(b[None, ...])
+ else:
+ batch2 = []
+ for b in batch:
+ batch2.append(b[None, ...])
+ return batch2
+ else:
+ raise NotImplementedError
diff --git a/models/grit_src/grit/modeling/text/modeling_bert.py b/models/grit_src/grit/modeling/text/modeling_bert.py
new file mode 100644
index 0000000000000000000000000000000000000000..3f8bf2d5d7552ee6c314da86a19a56eb0bdaa03e
--- /dev/null
+++ b/models/grit_src/grit/modeling/text/modeling_bert.py
@@ -0,0 +1,529 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch BERT model. """
+# Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+import copy
+import os
+import json
+import logging
+import math
+import sys
+from io import open
+import torch
+from torch import nn
+import torch.utils.checkpoint as checkpoint
+from .file_utils import cached_path
+
+
+logger = logging.getLogger()
+
+
+BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+ 'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json",
+ 'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-config.json",
+ 'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-config.json",
+ 'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-config.json",
+ 'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-config.json",
+ 'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-config.json",
+ 'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-config.json",
+ 'bert-base-german-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-config.json",
+ 'bert-large-uncased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-config.json",
+ 'bert-large-cased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-config.json",
+ 'bert-large-uncased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-config.json",
+ 'bert-large-cased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-config.json",
+ 'bert-base-cased-finetuned-mrpc': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-config.json",
+}
+
+
+def qk2attn(query, key, attention_mask, gamma):
+ query = query / gamma
+ attention_scores = torch.matmul(query, key.transpose(-1, -2))
+ if attention_mask is not None:
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+ attention_scores = attention_scores + attention_mask
+ return attention_scores.softmax(dim=-1)
+
+
+class QK2Attention(nn.Module):
+ def forward(self, query, key, attention_mask, gamma):
+ return qk2attn(query, key, attention_mask, gamma)
+
+
+LayerNormClass = torch.nn.LayerNorm
+
+
+class BertSelfAttention(nn.Module):
+ def __init__(self, config):
+ super(BertSelfAttention, self).__init__()
+ if config.hidden_size % config.num_attention_heads != 0:
+ raise ValueError(
+ "The hidden size (%d) is not a multiple of the number of attention "
+ "heads (%d)" % (config.hidden_size, config.num_attention_heads))
+ self.output_attentions = config.output_attentions
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+ self.softmax = nn.Softmax(dim=-1)
+ self.qk2attn = QK2Attention()
+
+ def transpose_for_scores(self, x):
+ if torch._C._get_tracing_state():
+ # exporter is not smart enough to detect dynamic size for some paths
+ x = x.view(x.shape[0], -1, self.num_attention_heads, self.attention_head_size)
+ else:
+ new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+ x = x.view(*new_x_shape)
+ return x.permute(0, 2, 1, 3)
+
+ def forward(self, hidden_states, attention_mask, head_mask=None,
+ history_state=None):
+ if history_state is not None:
+ x_states = torch.cat([history_state, hidden_states], dim=1)
+ mixed_query_layer = self.query(hidden_states)
+ mixed_key_layer = self.key(x_states)
+ mixed_value_layer = self.value(x_states)
+ else:
+ mixed_query_layer = self.query(hidden_states)
+ mixed_key_layer = self.key(hidden_states)
+ mixed_value_layer = self.value(hidden_states)
+
+ query_layer = self.transpose_for_scores(mixed_query_layer)
+ key_layer = self.transpose_for_scores(mixed_key_layer)
+ value_layer = self.transpose_for_scores(mixed_value_layer)
+
+ attention_probs = self.qk2attn(query_layer, key_layer, attention_mask, math.sqrt(self.attention_head_size))
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+
+ # Mask heads if we want to
+ if head_mask is not None:
+ attention_probs = attention_probs * head_mask
+
+ context_layer = torch.matmul(attention_probs, value_layer)
+
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(*new_context_layer_shape)
+
+ outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
+ return outputs
+
+
+class BertSelfOutput(nn.Module):
+ def __init__(self, config):
+ super(BertSelfOutput, self).__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.pre_norm = hasattr(config, 'pre_norm') and config.pre_norm
+ if not self.pre_norm:
+ self.LayerNorm = LayerNormClass(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states, input_tensor):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ if not self.pre_norm:
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ else:
+ hidden_states = hidden_states + input_tensor
+ return hidden_states
+
+
+class BertAttention(nn.Module):
+ def __init__(self, config):
+ super(BertAttention, self).__init__()
+ self.pre_norm = hasattr(config, 'pre_norm') and config.pre_norm
+ if self.pre_norm:
+ self.LayerNorm = LayerNormClass(config.hidden_size, eps=config.layer_norm_eps)
+ self.self = BertSelfAttention(config)
+ self.output = BertSelfOutput(config)
+
+ def forward(self, input_tensor, attention_mask, head_mask=None,
+ history_state=None):
+ if self.pre_norm:
+ self_outputs = self.self(self.LayerNorm(input_tensor), attention_mask, head_mask,
+ self.layerNorm(history_state) if history_state else history_state)
+ else:
+ self_outputs = self.self(input_tensor, attention_mask, head_mask,
+ history_state)
+ attention_output = self.output(self_outputs[0], input_tensor)
+ outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+class BertIntermediate(nn.Module):
+ def __init__(self, config):
+ super(BertIntermediate, self).__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ assert config.hidden_act == 'gelu', 'Please implement other activation functions'
+ self.intermediate_act_fn = _gelu_python
+
+ def forward(self, hidden_states):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+class BertOutput(nn.Module):
+ def __init__(self, config):
+ super(BertOutput, self).__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.pre_norm = hasattr(config, 'pre_norm') and config.pre_norm
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ if not self.pre_norm:
+ self.LayerNorm = LayerNormClass(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, hidden_states, input_tensor):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ if not self.pre_norm:
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ else:
+ hidden_states = hidden_states + input_tensor
+ return hidden_states
+
+
+class Mlp(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.pre_norm = hasattr(config, 'pre_norm') and config.pre_norm
+ self.intermediate = BertIntermediate(config)
+ if self.pre_norm:
+ self.LayerNorm = LayerNormClass(config.hidden_size, eps=config.layer_norm_eps)
+ self.output = BertOutput(config)
+
+ def forward(self, attention_output):
+ if not self.pre_norm:
+ intermediate_output = self.intermediate(attention_output)
+ else:
+ intermediate_output = self.intermediate(self.LayerNorm(attention_output))
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+class BertLayer(nn.Module):
+ def __init__(self, config, use_act_checkpoint=True):
+ super(BertLayer, self).__init__()
+ self.pre_norm = hasattr(config, 'pre_norm') and config.pre_norm
+ self.use_mlp_wrapper = hasattr(config, 'use_mlp_wrapper') and config.use_mlp_wrapper
+ self.attention = BertAttention(config)
+ self.use_act_checkpoint = use_act_checkpoint
+ if self.use_mlp_wrapper:
+ self.mlp = Mlp(config)
+ else:
+ self.intermediate = BertIntermediate(config)
+ if self.pre_norm:
+ self.LayerNorm = LayerNormClass(config.hidden_size, eps=config.layer_norm_eps)
+ self.output = BertOutput(config)
+
+ def forward(self, hidden_states, attention_mask, head_mask=None,
+ history_state=None):
+ if self.use_act_checkpoint:
+ attention_outputs = checkpoint.checkpoint(self.attention, hidden_states,
+ attention_mask, head_mask, history_state)
+ else:
+ attention_outputs = self.attention(hidden_states, attention_mask,
+ head_mask, history_state)
+ attention_output = attention_outputs[0]
+ if self.use_mlp_wrapper:
+ layer_output = self.mlp(attention_output)
+ else:
+ if not self.pre_norm:
+ intermediate_output = self.intermediate(attention_output)
+ else:
+ intermediate_output = self.intermediate(self.LayerNorm(attention_output))
+ layer_output = self.output(intermediate_output, attention_output)
+ outputs = (layer_output,) + attention_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+class BertEncoder(nn.Module):
+ def __init__(self, config, use_act_checkpoint=True):
+ super(BertEncoder, self).__init__()
+ self.output_attentions = config.output_attentions
+ self.output_hidden_states = config.output_hidden_states
+ self.layer = nn.ModuleList([BertLayer(config, use_act_checkpoint=use_act_checkpoint) for _ in range(config.num_hidden_layers)])
+ self.pre_norm = hasattr(config, 'pre_norm') and config.pre_norm
+ if self.pre_norm:
+ self.LayerNorm = LayerNormClass(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, hidden_states, attention_mask, head_mask=None,
+ encoder_history_states=None):
+ all_hidden_states = ()
+ all_attentions = ()
+ for i, layer_module in enumerate(self.layer):
+ if self.output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ history_state = None if encoder_history_states is None else encoder_history_states[i]
+ layer_outputs = layer_module(
+ hidden_states, attention_mask,
+ (None if head_mask is None else head_mask[i]),
+ history_state,
+ )
+ hidden_states = layer_outputs[0]
+
+ if self.output_attentions:
+ all_attentions = all_attentions + (layer_outputs[1],)
+ if self.pre_norm:
+ hidden_states = self.LayerNorm(hidden_states)
+ outputs = (hidden_states,)
+ if self.output_hidden_states:
+ outputs = outputs + (all_hidden_states,)
+ if self.output_attentions:
+ outputs = outputs + (all_attentions,)
+ return outputs
+
+CONFIG_NAME = "config.json"
+
+class PretrainedConfig(object):
+ """ Base class for all configuration classes.
+ Handle a few common parameters and methods for loading/downloading/saving configurations.
+ """
+ pretrained_config_archive_map = {}
+
+ def __init__(self, **kwargs):
+ self.finetuning_task = kwargs.pop('finetuning_task', None)
+ self.num_labels = kwargs.pop('num_labels', 2)
+ self.output_attentions = kwargs.pop('output_attentions', False)
+ self.output_hidden_states = kwargs.pop('output_hidden_states', False)
+ self.torchscript = kwargs.pop('torchscript', False)
+
+ def save_pretrained(self, save_directory):
+ """ Save a configuration object to a directory, so that it
+ can be re-loaded using the `from_pretrained(save_directory)` class method.
+ """
+ assert os.path.isdir(save_directory), "Saving path should be a directory where the model and configuration can be saved"
+
+ # If we save using the predefined names, we can load using `from_pretrained`
+ output_config_file = os.path.join(save_directory, CONFIG_NAME)
+
+ self.to_json_file(output_config_file)
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+ r""" Instantiate a PretrainedConfig from a pre-trained model configuration.
+
+ Params:
+ **pretrained_model_name_or_path**: either:
+ - a string with the `shortcut name` of a pre-trained model configuration to load from cache
+ or download and cache if not already stored in cache (e.g. 'bert-base-uncased').
+ - a path to a `directory` containing a configuration file saved
+ using the `save_pretrained(save_directory)` method.
+ - a path or url to a saved configuration `file`.
+ **cache_dir**: (`optional`) string:
+ Path to a directory in which a downloaded pre-trained model
+ configuration should be cached if the standard cache should not be used.
+ **return_unused_kwargs**: (`optional`) bool:
+ - If False, then this function returns just the final configuration object.
+ - If True, then this functions returns a tuple `(config, unused_kwargs)` where `unused_kwargs`
+ is a dictionary consisting of the key/value pairs whose keys are not configuration attributes:
+ ie the part of kwargs which has not been used to update `config` and is otherwise ignored.
+ **kwargs**: (`optional`) dict:
+ Dictionary of key/value pairs with which to update the configuration object after loading.
+ - The values in kwargs of any keys which are configuration attributes will be used
+ to override the loaded values.
+ - Behavior concerning key/value pairs whose keys are *not* configuration attributes is controlled
+ by the `return_unused_kwargs` keyword parameter.
+
+ Examples::
+
+ >>> config = BertConfig.from_pretrained('bert-base-uncased') # Download configuration from S3 and cache.
+ >>> config = BertConfig.from_pretrained('./test/saved_model/') # E.g. config (or model) was saved using `save_pretrained('./test/saved_model/')`
+ >>> config = BertConfig.from_pretrained('./test/saved_model/my_configuration.json')
+ >>> config = BertConfig.from_pretrained('bert-base-uncased', output_attention=True, foo=False)
+ >>> assert config.output_attention == True
+ >>> config, unused_kwargs = BertConfig.from_pretrained('bert-base-uncased', output_attention=True,
+ >>> foo=False, return_unused_kwargs=True)
+ >>> assert config.output_attention == True
+ >>> assert unused_kwargs == {'foo': False}
+
+ """
+ cache_dir = kwargs.pop('cache_dir', None)
+ return_unused_kwargs = kwargs.pop('return_unused_kwargs', False)
+
+ if pretrained_model_name_or_path in cls.pretrained_config_archive_map:
+ config_file = cls.pretrained_config_archive_map[pretrained_model_name_or_path]
+ elif os.path.isdir(pretrained_model_name_or_path):
+ config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
+ else:
+ config_file = pretrained_model_name_or_path
+ # redirect to the cache, if necessary
+ try:
+ resolved_config_file = cached_path(config_file, cache_dir=cache_dir)
+ except EnvironmentError:
+ if pretrained_model_name_or_path in cls.pretrained_config_archive_map:
+ logger.error(
+ "Couldn't reach server at '{}' to download pretrained model configuration file.".format(
+ config_file))
+ else:
+ logger.error(
+ "Model name '{}' was not found in model name list ({}). "
+ "We assumed '{}' was a path or url but couldn't find any file "
+ "associated to this path or url.".format(
+ pretrained_model_name_or_path,
+ ', '.join(cls.pretrained_config_archive_map.keys()),
+ config_file))
+ return None
+ if resolved_config_file == config_file:
+ logger.info("loading configuration file {}".format(config_file))
+ else:
+ logger.info("loading configuration file {} from cache at {}".format(
+ config_file, resolved_config_file))
+
+ # Load config
+ config = cls.from_json_file(resolved_config_file)
+
+ # Update config with kwargs if needed
+ to_remove = []
+ for key, value in kwargs.items():
+ if hasattr(config, key):
+ setattr(config, key, value)
+ to_remove.append(key)
+ # add img_layer_norm_eps, use_img_layernorm
+ if "img_layer_norm_eps" in kwargs:
+ setattr(config, "img_layer_norm_eps", kwargs["img_layer_norm_eps"])
+ to_remove.append("img_layer_norm_eps")
+ if "use_img_layernorm" in kwargs:
+ setattr(config, "use_img_layernorm", kwargs["use_img_layernorm"])
+ to_remove.append("use_img_layernorm")
+ for key in to_remove:
+ kwargs.pop(key, None)
+
+ logger.info("Model config %s", config)
+ if return_unused_kwargs:
+ return config, kwargs
+ else:
+ return config
+
+ @classmethod
+ def from_dict(cls, json_object):
+ """Constructs a `Config` from a Python dictionary of parameters."""
+ config = cls(vocab_size_or_config_json_file=-1)
+ for key, value in json_object.items():
+ config.__dict__[key] = value
+ return config
+
+ @classmethod
+ def from_json_file(cls, json_file):
+ """Constructs a `BertConfig` from a json file of parameters."""
+ with open(json_file, "r", encoding='utf-8') as reader:
+ text = reader.read()
+ return cls.from_dict(json.loads(text))
+
+ def __eq__(self, other):
+ return self.__dict__ == other.__dict__
+
+ def __repr__(self):
+ return str(self.to_json_string())
+
+ def to_dict(self):
+ """Serializes this instance to a Python dictionary."""
+ output = copy.deepcopy(self.__dict__)
+ return output
+
+ def to_json_string(self):
+ """Serializes this instance to a JSON string."""
+ return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+
+ def to_json_file(self, json_file_path):
+ """ Save this instance to a json file."""
+ with open(json_file_path, "w", encoding='utf-8') as writer:
+ writer.write(self.to_json_string())
+
+
+class BertConfig(PretrainedConfig):
+ r"""
+ :class:`~pytorch_transformers.BertConfig` is the configuration class to store the configuration of a
+ `BertModel`.
+
+
+ Arguments:
+ vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`.
+ hidden_size: Size of the encoder layers and the pooler layer.
+ num_hidden_layers: Number of hidden layers in the Transformer encoder.
+ num_attention_heads: Number of attention heads for each attention layer in
+ the Transformer encoder.
+ intermediate_size: The size of the "intermediate" (i.e., feed-forward)
+ layer in the Transformer encoder.
+ hidden_act: The non-linear activation function (function or string) in the
+ encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
+ hidden_dropout_prob: The dropout probabilitiy for all fully connected
+ layers in the embeddings, encoder, and pooler.
+ attention_probs_dropout_prob: The dropout ratio for the attention
+ probabilities.
+ max_position_embeddings: The maximum sequence length that this model might
+ ever be used with. Typically set this to something large just in case
+ (e.g., 512 or 1024 or 2048).
+ type_vocab_size: The vocabulary size of the `token_type_ids` passed into
+ `BertModel`.
+ initializer_range: The sttdev of the truncated_normal_initializer for
+ initializing all weight matrices.
+ layer_norm_eps: The epsilon used by LayerNorm.
+ """
+ pretrained_config_archive_map = BERT_PRETRAINED_CONFIG_ARCHIVE_MAP
+
+ def __init__(self,
+ vocab_size_or_config_json_file=30522,
+ hidden_size=768,
+ num_hidden_layers=12,
+ num_attention_heads=12,
+ intermediate_size=3072,
+ hidden_act="gelu",
+ hidden_dropout_prob=0.1,
+ attention_probs_dropout_prob=0.1,
+ max_position_embeddings=512,
+ type_vocab_size=2,
+ initializer_range=0.02,
+ layer_norm_eps=1e-12,
+ **kwargs):
+ super(BertConfig, self).__init__(**kwargs)
+ if isinstance(vocab_size_or_config_json_file, str):
+ with open(vocab_size_or_config_json_file, "r", encoding='utf-8') as reader:
+ json_config = json.loads(reader.read())
+ for key, value in json_config.items():
+ self.__dict__[key] = value
+ elif isinstance(vocab_size_or_config_json_file, int):
+ self.vocab_size = vocab_size_or_config_json_file
+ self.hidden_size = hidden_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.hidden_act = hidden_act
+ self.intermediate_size = intermediate_size
+ self.hidden_dropout_prob = hidden_dropout_prob
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
+ self.max_position_embeddings = max_position_embeddings
+ self.type_vocab_size = type_vocab_size
+ self.initializer_range = initializer_range
+ self.layer_norm_eps = layer_norm_eps
+ else:
+ raise ValueError("First argument must be either a vocabulary size (int)"
+ "or the path to a pretrained model config file (str)")
+
+
+def _gelu_python(x):
+
+ return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
\ No newline at end of file
diff --git a/models/grit_src/grit/modeling/text/text_decoder.py b/models/grit_src/grit/modeling/text/text_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..071baa7a52d21d7132cc492f070cba066d17aa43
--- /dev/null
+++ b/models/grit_src/grit/modeling/text/text_decoder.py
@@ -0,0 +1,672 @@
+# Modified by Jialian Wu from
+# https://github.com/microsoft/GenerativeImage2Text/blob/main/generativeimage2text/layers/decoder.py
+# and https://github.com/kdexd/virtex
+from torch import nn
+import torch
+import functools
+from torch.nn import functional as F
+import warnings
+
+
+class TextualHead(nn.Module):
+ def __init__(self,
+ visual_feature_size: int, vocab_size: int, hidden_size: int):
+ super().__init__()
+ self.visual_feature_size = visual_feature_size
+ self.vocab_size = vocab_size
+ self.hidden_size = hidden_size
+
+ @property
+ def textual_feature_size(self):
+ return self.hidden_size
+
+
+class WordAndPositionalEmbedding(nn.Module):
+ def __init__(
+ self,
+ vocab_size: int,
+ hidden_size: int,
+ dropout: float = 0.0,
+ max_caption_length: int = 30,
+ padding_idx: int = 0,
+ ):
+ super().__init__()
+ self.vocab_size = vocab_size
+ self.padding_idx = padding_idx
+
+ #self.words = nn.Embedding(vocab_size, hidden_size, padding_idx=padding_idx)
+ self.words = nn.Embedding(vocab_size, hidden_size)
+
+ # We provide no "padding index" for positional embeddings. We zero out
+ # the positional embeddings of padded positions as a post-processing.
+ self.positions = nn.Embedding(max_caption_length, hidden_size)
+ self.layer_norm = nn.LayerNorm(
+ hidden_size, eps=1e-8, elementwise_affine=True
+ )
+ self.dropout = nn.Dropout(p=dropout)
+
+ def forward(self, tokens: torch.Tensor):
+ position_indices = self._create_position_indices(tokens)
+
+ # shape: (batch_size, max_caption_length, hidden_size)
+ word_embeddings = self.words(tokens)
+ position_embeddings = self.positions(position_indices)
+
+ # shape: (batch_size, max_caption_length, hidden_size)
+ embeddings = self.layer_norm(word_embeddings + position_embeddings)
+ embeddings = self.dropout(embeddings)
+
+ return embeddings
+
+ @functools.lru_cache(maxsize=128)
+ def _create_position_indices(self, tokens: torch.Tensor):
+
+ # Create position indices of the same size as token indices.
+ batch_size, max_caption_length = tokens.size()
+ positions = torch.arange(
+ max_caption_length, dtype=tokens.dtype, device=tokens.device
+ )
+ # shape: (batch_size, max_caption_length)
+ positions = positions.unsqueeze(0).expand(batch_size, max_caption_length)
+ return positions
+
+
+class BertEncoderAsDecoder(nn.Module):
+ def __init__(self, encoder):
+ super().__init__()
+ self.encoder = encoder
+
+ def forward(self, tgt, memory,
+ tgt_mask=None,
+ tgt_key_padding_mask=None,
+ memory_key_padding_mask=None,
+ tgt_bi_valid_mask=None,
+ encoder_history_states=None,
+ ):
+ assert tgt_key_padding_mask is None, 'not supported'
+ assert tgt_mask.dim() == 2
+ assert tgt_mask.shape[0] == tgt_mask.shape[1]
+ # tgt_mask should always be 0/negative infinity
+ tgt = tgt.transpose(0, 1)
+ memory = memory.transpose(0, 1)
+
+ hidden_states = torch.cat((memory, tgt), dim=1)
+ num_tgt = tgt.shape[1]
+ num_memory = memory.shape[1]
+ device = tgt.device
+ dtype = tgt.dtype
+ top_left = torch.zeros((num_memory, num_memory), device=device, dtype=dtype)
+ top_right = torch.full((num_memory, num_tgt), float('-inf'), device=tgt.device, dtype=dtype,)
+ bottom_left = torch.zeros((num_tgt, num_memory), dtype=dtype, device=tgt_mask.device,)
+ left = torch.cat((top_left, bottom_left), dim=0)
+ right = torch.cat((top_right, tgt_mask.to(dtype)), dim=0)
+
+ full_attention_mask = torch.cat((left, right), dim=1)[None, :]
+
+ if memory_key_padding_mask is None:
+ memory_key_padding_mask = torch.full((memory.shape[0], memory.shape[1]), fill_value=False, device=device)
+ # if it is False, it means valid. That is, it is not a padding
+ assert memory_key_padding_mask.dtype == torch.bool
+ zero_negative_infinity = torch.zeros_like(memory_key_padding_mask, dtype=tgt.dtype)
+ zero_negative_infinity[memory_key_padding_mask] = float('-inf')
+ full_attention_mask = full_attention_mask.expand((memory_key_padding_mask.shape[0], num_memory + num_tgt, num_memory + num_tgt))
+ full_attention_mask = full_attention_mask.clone()
+ origin_left = full_attention_mask[:, :, :num_memory]
+ update = zero_negative_infinity[:, None, :]
+ full_attention_mask[:, :, :num_memory] = origin_left + update
+
+ if tgt_bi_valid_mask is not None:
+ # verify the correctness
+ bs = full_attention_mask.shape[0]
+ # during inference, tgt_bi_valid_mask's length is not changed, but
+ # num_tgt can be increased
+ max_valid_target = tgt_bi_valid_mask.shape[1]
+ mask = tgt_bi_valid_mask[:, None, :].expand((bs, num_memory+num_tgt, max_valid_target))
+ full_attention_mask[:, :, num_memory:(num_memory+max_valid_target)][mask] = 0
+
+ # add axis for multi-head
+ full_attention_mask = full_attention_mask[:, None, :, :]
+
+ if encoder_history_states is None:
+ result = self.encoder(
+ hidden_states=hidden_states,
+ attention_mask=full_attention_mask,
+ encoder_history_states=encoder_history_states,
+ )
+ result = list(result)
+ result[0] = result[0][:, num_memory:].transpose(0, 1)
+ if self.encoder.output_hidden_states:
+ return result[0], result[1]
+ else:
+ # make it back-compatible
+ return result[0]
+ else:
+ encoder_out = self.encoder(
+ hidden_states=hidden_states[:, -1:],
+ attention_mask=full_attention_mask[:, :, -1:],
+ encoder_history_states=encoder_history_states,
+ )
+ result = encoder_out[0].transpose(0, 1)
+ if self.encoder.output_hidden_states:
+ return result, encoder_out[1]
+ else:
+ return result
+
+
+def create_transformer(decoder_type, norm_type,
+ textual_feature_size,
+ attention_heads,
+ feedforward_size,
+ dropout,
+ num_layers,
+ output_hidden_states=False,
+ use_mlp_wrapper=None,
+ use_act_checkpoint=True,
+ ):
+ assert norm_type in ['post', 'pre']
+ if decoder_type is None:
+ LayerClass = (
+ nn.TransformerDecoderLayer
+ if norm_type == "post"
+ else PreNormTransformerDecoderLayer
+ )
+ _layer = LayerClass(
+ textual_feature_size,
+ attention_heads,
+ dim_feedforward=feedforward_size,
+ dropout=dropout,
+ activation="gelu",
+ )
+ return nn.TransformerDecoder(_layer, num_layers)
+ elif decoder_type == 'bert_en':
+ from .modeling_bert import BertConfig, BertEncoder
+ config = BertConfig(
+ vocab_size_or_config_json_file=30522,
+ hidden_size=textual_feature_size,
+ num_hidden_layers=num_layers,
+ num_attention_heads=attention_heads,
+ intermediate_size=feedforward_size,
+ hidden_act="gelu",
+ hidden_dropout_prob=0.1,
+ attention_probs_dropout_prob=0.1,
+ layer_norm_eps=1e-12,
+ )
+ config.pre_norm = (norm_type == 'pre')
+ config.use_mlp_wrapper = use_mlp_wrapper
+ config.output_hidden_states = output_hidden_states
+ encoder = BertEncoder(config, use_act_checkpoint=use_act_checkpoint)
+ return BertEncoderAsDecoder(encoder)
+
+
+class PreNormTransformerDecoderLayer(nn.TransformerDecoderLayer):
+ def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
+ tgt_key_padding_mask=None, memory_key_padding_mask=None):
+ # fmt: off
+ # We use the members (modules) from super-class, just the order of
+ # operations is changed here. First layernorm, then attention.
+ tgt2 = self.norm1(tgt)
+ tgt2, _ = self.self_attn(
+ tgt2, tgt2, tgt2, attn_mask=tgt_mask,
+ key_padding_mask=tgt_key_padding_mask
+ )
+ tgt = tgt + self.dropout1(tgt2)
+
+ # Layernorm first, then decoder attention.
+ tgt2 = self.norm2(tgt)
+ tgt2, _ = self.multihead_attn(
+ tgt2, memory, memory, attn_mask=memory_mask,
+ key_padding_mask=memory_key_padding_mask
+ )
+ tgt = tgt + self.dropout2(tgt2)
+
+ # Layernorm first, then transformation through feedforward network.
+ tgt2 = self.norm3(tgt)
+ tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+ tgt = tgt + self.dropout3(tgt2)
+ return tgt
+
+
+class TransformerDecoderTextualHead(TextualHead):
+ def __init__(
+ self,
+ object_feature_size: int,
+ vocab_size: int,
+ hidden_size: int,
+ num_layers: int,
+ attention_heads: int,
+ feedforward_size: int,
+ dropout: float = 0.1,
+ norm_type: str = "post",
+ mask_future_positions: bool = True,
+ max_caption_length: int = 1024,
+ padding_idx: int = 0,
+ decoder_type=None,
+ not_tie_weight=None,
+ output_hidden_states=None,
+ use_mlp_wrapper=None,
+ use_act_checkpoint=True,
+ ):
+ super().__init__(object_feature_size, vocab_size, hidden_size)
+ self.num_layers = num_layers
+ self.attention_heads = attention_heads
+ self.feedforward_size = feedforward_size
+ self.dropout = dropout
+ assert mask_future_positions
+ self.padding_idx = padding_idx
+
+ self.object_feature_projection = nn.Sequential(
+ nn.Linear(object_feature_size, self.textual_feature_size),
+ nn.LayerNorm(self.textual_feature_size))
+
+ self.embedding = WordAndPositionalEmbedding(
+ self.vocab_size,
+ self.textual_feature_size,
+ dropout=dropout,
+ max_caption_length=max_caption_length,
+ padding_idx=padding_idx,
+ )
+ self.transformer = create_transformer(
+ decoder_type=decoder_type,
+ norm_type=norm_type,
+ textual_feature_size=self.textual_feature_size,
+ attention_heads=self.attention_heads,
+ feedforward_size=self.feedforward_size,
+ dropout=dropout,
+ num_layers=self.num_layers,
+ output_hidden_states=output_hidden_states,
+ use_mlp_wrapper=use_mlp_wrapper,
+ use_act_checkpoint=use_act_checkpoint,
+ )
+ self.apply(self._init_weights)
+
+ # Create an output linear layer and tie the input and output word
+ # embeddings to reduce parametejs.
+ self.output = nn.Linear(self.textual_feature_size, vocab_size)
+ if not not_tie_weight:
+ self.output.weight = self.embedding.words.weight
+
+ @staticmethod
+ def _init_weights(module):
+ """Initialize weights like BERT - N(0.0, 0.02), bias = 0."""
+
+ if isinstance(module, nn.Linear):
+ module.weight.data.normal_(mean=0.0, std=0.02)
+ elif isinstance(module, nn.MultiheadAttention):
+ module.in_proj_weight.data.normal_(mean=0.0, std=0.02)
+ module.out_proj.weight.data.normal_(mean=0.0, std=0.02)
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=0.02)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+ def forward(
+ self,
+ hidden_states,
+ text_tokens,
+ ):
+ projected_object_features = self.object_feature_projection(hidden_states) if hidden_states is not None else None
+ batch_size, max_text_length = text_tokens.size()
+ text_embeddings = self.embedding(text_tokens)
+
+ # An additive mask for masking the future (one direction).
+ uni_mask_zero_neg = self._generate_future_mask(
+ max_text_length, text_embeddings.dtype, text_embeddings.device
+ )
+
+ # We transpose the first two dimensions of tokens embeddings and visual
+ # features, as required by decoder.
+ text_embeddings = text_embeddings.transpose(0, 1)
+
+ projected_object_features = projected_object_features.transpose(0, 1)
+
+ # if transformer here is the pytorch/decoder, there is no chance, the
+ # output is always tensor
+ trans_out = self.transformer(
+ text_embeddings,
+ projected_object_features,
+ tgt_mask=uni_mask_zero_neg,
+ )
+ if isinstance(trans_out, tuple):
+ textual_features = trans_out[0]
+ else:
+ assert isinstance(trans_out, torch.Tensor)
+ textual_features = trans_out
+ # Undo the transpose and bring batch to dim 0.
+ # shape: (batch_size, max_caption_length, hidden_size)
+ textual_features = textual_features.transpose(0, 1)
+
+ # shape: (batch_size, max_caption_length, vocab_size)
+ output_logits = self.output(textual_features)
+ if isinstance(trans_out, tuple):
+ return output_logits, trans_out[1]
+ else:
+ return output_logits
+
+ def _generate_future_mask(
+ self, size: int, dtype: torch.dtype, device: torch.device
+ ):
+ # Default mask is for forward direction. Flip for backward direction.
+ mask = torch.triu(
+ torch.ones(size, size, device=device, dtype=dtype), diagonal=1
+ )
+ mask = mask.masked_fill(mask == 1, float("-inf"))
+ return mask
+
+
+class AutoRegressiveBeamSearch(object):
+ def __init__(
+ self,
+ end_token_id: int,
+ max_steps: int = 50,
+ beam_size: int = 5,
+ objectdet=True,
+ per_node_beam_size: int = 2,
+ ):
+ self._eos_index = end_token_id
+ self.max_steps = max_steps
+ self.beam_size = beam_size
+ self.objectdet = objectdet
+ self.per_node_beam_size = per_node_beam_size or beam_size
+
+ def search(self, begin_tokens, step):
+ if self.beam_size > 1 and self.objectdet:
+ only_return_best = False
+ else:
+ only_return_best = True
+
+ batch_size = begin_tokens.size()[0]
+
+ predictions = begin_tokens.unsqueeze(1).expand((batch_size, self.beam_size, begin_tokens.shape[-1]))
+ # Calculate the first timestep. This is done outside the main loop
+ # because we are going from a single decoder input (the output from the
+ # encoder) to the top `beam_size` decoder outputs. On the other hand,
+ # within the main loop we are going from the `beam_size` elements of the
+ # beam to `beam_size`^2 candidates from which we will select the top
+ # `beam_size` elements for the next iteration.
+ # shape: (batch_size, num_classes)
+ start_class_logits = step(begin_tokens)
+
+ # Convert logits to logprobs.
+ # shape: (batch_size * beam_size, vocab_size)
+ start_class_logprobs = F.log_softmax(start_class_logits, dim=1)
+
+ num_classes = start_class_logprobs.size()[1]
+
+ # shape: (batch_size, beam_size), (batch_size, beam_size)
+ start_top_logprobs, start_predicted_classes = start_class_logprobs.topk(
+ self.beam_size
+ )
+
+ if (
+ self.beam_size == 1
+ and (start_predicted_classes == self._eos_index).all()
+ ):
+ warnings.warn(
+ "Empty object description predicted. You may want to increase beam"
+ "size or ensure your step function is working properly.",
+ RuntimeWarning,
+ )
+ if only_return_best:
+ return start_predicted_classes, start_top_logprobs
+ else:
+ return start_predicted_classes.unsqueeze(-1), start_top_logprobs
+
+ # The log probs for the last time step.
+ # shape: (batch_size, beam_size)
+ last_logprobs = start_top_logprobs
+
+ # shape: (batch_size, beam_size, sequence_length)
+ predictions = torch.cat([predictions, start_predicted_classes.unsqueeze(-1)], dim=-1)
+
+ # Log probability tensor that mandates that the end token is selected.
+ # shape: (batch_size * beam_size, num_classes)
+ logprobs_after_end = start_class_logprobs.new_full(
+ (batch_size * self.beam_size, num_classes), float("-inf")
+ )
+ logprobs_after_end[:, self._eos_index] = 0.0
+
+ logits_after_end = start_class_logprobs.new_full(
+ (batch_size * self.beam_size, num_classes), float("-inf")
+ )
+ logits_after_end[:, self._eos_index] = 0
+
+ while predictions.shape[-1] < self.max_steps:
+ # shape: (batch_size * beam_size,)
+ last_predictions = predictions[:, :, -1].reshape(batch_size * self.beam_size)
+
+ # If every predicted token from the last step is `self._eos_index`,
+ # then we can stop early.
+ if (last_predictions == self._eos_index).all():
+ break
+
+ predictions_so_far = predictions.view(
+ batch_size * self.beam_size, -1
+ )
+ # shape: (batch_size * beam_size, num_classes)
+ class_logits = step(predictions_so_far)
+
+ # Set logprobs of last predicted tokens as high negative value to avoid
+ # repetition in description.
+ class_logits = class_logits.scatter(1, predictions_so_far[:, -1].view((-1, 1)), -10000)
+
+ # shape: (batch_size * beam_size, num_classes)
+ last_predictions_expanded = last_predictions.unsqueeze(-1).expand(
+ batch_size * self.beam_size, num_classes
+ )
+
+ # Here we are finding any beams where we predicted the end token in
+ # the previous timestep and replacing the distribution with a
+ # one-hot distribution, forcing the beam to predict the end token
+ # this timestep as well.
+ class_logits = torch.where(
+ last_predictions_expanded == self._eos_index,
+ logits_after_end,
+ class_logits,
+ )
+
+ # Convert logits to logprobs.
+ # shape: (batch_size * beam_size, vocab_size)
+ class_logprobs = F.log_softmax(class_logits, dim=1)
+
+ # shape (both): (batch_size * beam_size, per_node_beam_size)
+ top_logprobs, predicted_classes = class_logprobs.topk(
+ self.per_node_beam_size
+ )
+
+ # Here we expand the last log probs to `(batch_size * beam_size,
+ # per_node_beam_size)` so that we can add them to the current log
+ # probs for this timestep. This lets us maintain the log
+ # probability of each element on the beam.
+ # shape: (batch_size * beam_size, per_node_beam_size)
+ expanded_last_logprobs = (
+ last_logprobs.unsqueeze(2)
+ .expand(batch_size, self.beam_size, self.per_node_beam_size)
+ .reshape(batch_size * self.beam_size, self.per_node_beam_size)
+ )
+ # shape: (batch_size * beam_size, per_node_beam_size)
+ summed_top_logprobs = top_logprobs + expanded_last_logprobs
+
+ # shape: (batch_size, beam_size * per_node_beam_size)
+ reshaped_summed = summed_top_logprobs.reshape(
+ batch_size, self.beam_size * self.per_node_beam_size
+ )
+ # shape: (batch_size, beam_size * per_node_beam_size)
+ reshaped_predicted_classes = predicted_classes.reshape(
+ batch_size, self.beam_size * self.per_node_beam_size
+ )
+ # Append the predictions to the current beam.
+ reshaped_beam = (
+ predictions.view(batch_size * self.beam_size, 1, -1)
+ .repeat(1, self.per_node_beam_size, 1)
+ .reshape(batch_size, self.beam_size * self.per_node_beam_size, -1)
+ )
+ # batch_size, (beam_size * per_node_beach_size), #token
+ reshaped_beam = torch.cat([reshaped_beam, reshaped_predicted_classes.unsqueeze(-1)], dim=-1)
+
+ # Keep only the top `beam_size` beam indices.
+ # shape: (batch_size, beam_size), (batch_size, beam_size)
+ restricted_beam_logprobs, restricted_beam_indices = reshaped_summed.topk(
+ self.beam_size
+ )
+ predictions = reshaped_beam.gather(
+ 1, restricted_beam_indices.unsqueeze(-1).repeat(1,1,reshaped_beam.shape[-1])
+ )
+
+ # shape: (batch_size, beam_size)
+ last_logprobs = restricted_beam_logprobs
+
+ if not torch.isfinite(last_logprobs).all():
+ warnings.warn(
+ "Infinite log probs encountered. Some final descriptions may not "
+ "make sense. This can happen when the beam size is larger than"
+ " the number of valid (non-zero probability) transitions that "
+ "the step function produces.",
+ RuntimeWarning,
+ )
+
+ # Optionally select best beam and its logprobs.
+ if only_return_best:
+ # shape: (batch_size, sequence_length)
+ predictions = predictions[:, 0, :]
+ last_logprobs = last_logprobs[:, 0]
+ num_valid = (predictions != self._eos_index).sum(dim=-1)
+ num_valid += (predictions == self._eos_index).sum(dim=-1) > 0
+ num_valid = num_valid - begin_tokens.shape[1]
+ num_valid = num_valid.clip(min=1)
+
+ last_logprobs = last_logprobs / num_valid
+
+ return predictions, last_logprobs
+
+
+class GRiTTextDecoder(nn.Module):
+ def __init__(
+ self,
+ transformer,
+ begin_token_id=101,
+ beamsearch_decode=None,
+ loss_type=None,
+ tokenizer=None,
+ ):
+ super().__init__()
+ self.textual = transformer
+ self.padding_idx = self.textual.padding_idx
+
+ self.begin_token_id = begin_token_id
+ self.beamsearch_decode = beamsearch_decode
+ self.tokenizer = tokenizer
+
+ if loss_type is None:
+ self.loss = nn.CrossEntropyLoss(ignore_index=self.padding_idx)
+ elif loss_type == 'smooth':
+ self.loss = SmoothLabelCrossEntropyLoss(ignore_index=self.padding_idx)
+ else:
+ raise NotImplementedError(loss_type)
+
+ def forward(self, batch):
+ object_features = batch['object_features']
+
+ if self.training:
+ caption_token_input = batch["text_tokens"]
+
+ output_logits = self.textual(
+ object_features,
+ caption_token_input,
+ )
+
+ if 'need_predict' in batch:
+ # in place should also be good, but we do not choose that for
+ # safety as we may use it in prediction results in future
+ target = batch["text_tokens"].clone()
+ target[batch['need_predict'] == 0] = self.padding_idx
+ else:
+ target = batch["text_tokens"]
+
+ feat = output_logits[:, :-1].contiguous()
+ target = target[:, 1:].contiguous()
+ feat = feat.view(-1, self.textual.vocab_size)
+ target = target.view(-1)
+
+ valid_mask = target != self.padding_idx
+ target = target[valid_mask]
+ feat = feat[valid_mask]
+ loss = self.loss(feat, target)
+
+ return loss
+ else:
+ output_dict = self.infer(object_features)
+ return output_dict
+
+ def infer(self, object_features):
+ batch_size = object_features.size(0)
+ begin_tokens = object_features.new_full(
+ (batch_size, 1), self.begin_token_id
+ ).long()
+
+ decoding_step = functools.partial(
+ self.decoding_step, object_features
+ )
+
+ object_description_tokens, logprobs = self.beamsearch_decode.search(
+ begin_tokens, decoding_step
+ )
+
+ output_dict = {
+ 'predictions': object_description_tokens,
+ 'logprobs': logprobs,
+ }
+
+ return output_dict
+
+ def decoding_step(self, object_features, partial_text):
+ batch_size = object_features.shape[0]
+ beam_size = int(partial_text.size(0) / batch_size)
+ if beam_size > 1:
+ batch_size, num_token, channels = object_features.size()
+ object_features = object_features.unsqueeze(1).repeat(1, beam_size, 1, 1)
+ object_features = object_features.view(
+ batch_size * beam_size, num_token, channels
+ )
+
+ text_lengths = torch.ones_like(partial_text)
+ if len(text_lengths.size()) != 2:
+ partial_text = partial_text.unsqueeze(1)
+
+ # shape: (batch_size * beam_size, partial_caption_length, vocab_size)
+ logits = self.textual(
+ object_features,
+ partial_text,
+ )
+
+ return logits[:, -1, :].float()
+
+
+class SmoothLabelCrossEntropyLoss(nn.Module):
+ def __init__(self, eps=0.1, log_prefix='', ignore_index=None):
+ super().__init__()
+ self.eps = eps
+ self.log_soft = nn.LogSoftmax(dim=1)
+ self.kl = nn.KLDivLoss(reduction='none')
+
+ self.iter = 0
+ self.max_loss = 0
+ self.min_loss = 0
+ self.log_prefix = log_prefix
+ self.ignore_index = ignore_index
+
+ def forward(self, feature, target):
+ feature = feature.float()
+ if self.ignore_index is not None:
+ valid_mask = target != self.ignore_index
+ target = target[valid_mask]
+ feature = feature[valid_mask]
+ assert target.numel() > 0
+ self.iter += 1
+ eps = self.eps
+ n_class = feature.size(1)
+ one_hot = torch.zeros_like(feature).scatter(1, target.view(-1, 1), 1)
+ one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
+ log_prb = self.log_soft(feature)
+ loss = self.kl(log_prb, one_hot)
+ return loss.sum(dim=1).mean()
+
diff --git a/models/grit_src/grit/predictor.py b/models/grit_src/grit/predictor.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c188ea2ab5fac232554d4eaaf2fb073670a70e4
--- /dev/null
+++ b/models/grit_src/grit/predictor.py
@@ -0,0 +1,66 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Jialian Wu from https://github.com/facebookresearch/detectron2/blob/main/detectron2/utils/visualizer.py
+import torch
+
+from detectron2.engine.defaults import DefaultPredictor
+from detectron2.utils.visualizer import ColorMode, Visualizer
+
+
+class Visualizer_GRiT(Visualizer):
+ def __init__(self, image, instance_mode=None):
+ super().__init__(image, instance_mode=instance_mode)
+
+ def draw_instance_predictions(self, predictions):
+ boxes = predictions.pred_boxes if predictions.has("pred_boxes") else None
+ scores = predictions.scores if predictions.has("scores") else None
+ classes = predictions.pred_classes.tolist() if predictions.has("pred_classes") else None
+ object_description = predictions.pred_object_descriptions.data
+ # uncomment to output scores in visualized images
+ # object_description = [c + '|' + str(round(s.item(), 1)) for c, s in zip(object_description, scores)]
+
+ if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get("thing_colors"):
+ colors = [
+ self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in classes
+ ]
+ alpha = 0.8
+ else:
+ colors = None
+ alpha = 0.5
+
+ if self._instance_mode == ColorMode.IMAGE_BW:
+ self.output.reset_image(
+ self._create_grayscale_image(
+ (predictions.pred_masks.any(dim=0) > 0).numpy()
+ if predictions.has("pred_masks")
+ else None
+ )
+ )
+ alpha = 0.3
+
+ self.overlay_instances(
+ masks=None,
+ boxes=boxes,
+ labels=object_description,
+ keypoints=None,
+ assigned_colors=colors,
+ alpha=alpha,
+ )
+ return self.output
+
+
+class VisualizationDemo(object):
+ def __init__(self, cfg, instance_mode=ColorMode.IMAGE):
+ self.cpu_device = torch.device("cpu")
+ self.instance_mode = instance_mode
+
+ self.predictor = DefaultPredictor(cfg)
+
+ def run_on_image(self, image):
+ predictions = self.predictor(image)
+ # Convert image from OpenCV BGR format to Matplotlib RGB format.
+ image = image[:, :, ::-1]
+ visualizer = Visualizer_GRiT(image, instance_mode=self.instance_mode)
+ instances = predictions["instances"].to(self.cpu_device)
+ vis_output = visualizer.draw_instance_predictions(predictions=instances)
+
+ return predictions, vis_output
\ No newline at end of file
diff --git a/models/grit_src/image_dense_captions.py b/models/grit_src/image_dense_captions.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fae11ca99f304aa682db5cc3d4ff295c6303182
--- /dev/null
+++ b/models/grit_src/image_dense_captions.py
@@ -0,0 +1,65 @@
+import argparse
+import multiprocessing as mp
+import os
+import time
+import cv2
+import tqdm
+import sys
+
+from detectron2.config import get_cfg
+from detectron2.data.detection_utils import read_image
+from detectron2.utils.logger import setup_logger
+
+sys.path.insert(0, 'models/grit_src/third_party/CenterNet2/projects/CenterNet2/')
+from centernet.config import add_centernet_config
+from models.grit_src.grit.config import add_grit_config
+
+from models.grit_src.grit.predictor import VisualizationDemo
+import json
+
+
+# constants
+WINDOW_NAME = "GRiT"
+
+
+def dense_pred_to_caption(predictions):
+ boxes = predictions["instances"].pred_boxes if predictions["instances"].has("pred_boxes") else None
+ object_description = predictions["instances"].pred_object_descriptions.data
+ new_caption = ""
+ for i in range(len(object_description)):
+ new_caption += (object_description[i] + ": " + str([int(a) for a in boxes[i].tensor.cpu().detach().numpy()[0]])) + "; "
+ return new_caption
+
+def setup_cfg(args):
+ cfg = get_cfg()
+ if args["cpu"]:
+ cfg.MODEL.DEVICE="cpu"
+ add_centernet_config(cfg)
+ add_grit_config(cfg)
+ cfg.merge_from_file(args["config_file"])
+ cfg.merge_from_list(args["opts"])
+ # Set score_threshold for builtin models
+ cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = args["confidence_threshold"]
+ cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = args["confidence_threshold"]
+ if args["test_task"]:
+ cfg.MODEL.TEST_TASK = args["test_task"]
+ cfg.MODEL.BEAM_SIZE = 1
+ cfg.MODEL.ROI_HEADS.SOFT_NMS_ENABLED = False
+ cfg.USE_ACT_CHECKPOINT = False
+ cfg.freeze()
+ return cfg
+
+
+def get_parser():
+ arg_dict = {'config_file': "models/grit_src/configs/GRiT_B_DenseCap_ObjectDet.yaml", 'cpu': False, 'confidence_threshold': 0.5, 'test_task': 'DenseCap', 'opts': ["MODEL.WEIGHTS", "pretrained_models/grit_b_densecap_objectdet.pth"]}
+ return arg_dict
+
+def image_caption_api(image_src):
+ args2 = get_parser()
+ cfg = setup_cfg(args2)
+ demo = VisualizationDemo(cfg)
+ if image_src:
+ img = read_image(image_src, format="BGR")
+ predictions, visualized_output = demo.run_on_image(img)
+ new_caption = dense_pred_to_caption(predictions)
+ return new_caption
\ No newline at end of file
diff --git a/models/grit_src/third_party/CenterNet2/.circleci/config.yml b/models/grit_src/third_party/CenterNet2/.circleci/config.yml
new file mode 100644
index 0000000000000000000000000000000000000000..097afade9aa0a48b54ab7c7aae93058227901e7a
--- /dev/null
+++ b/models/grit_src/third_party/CenterNet2/.circleci/config.yml
@@ -0,0 +1,256 @@
+version: 2.1
+
+# -------------------------------------------------------------------------------------
+# Environments to run the jobs in
+# -------------------------------------------------------------------------------------
+cpu: &cpu
+ machine:
+ image: ubuntu-2004:202107-02
+ resource_class: medium
+
+gpu: &gpu
+ machine:
+ # NOTE: use a cuda vesion that's supported by all our pytorch versions
+ image: ubuntu-1604-cuda-11.1:202012-01
+ resource_class: gpu.nvidia.small
+
+windows-cpu: &windows_cpu
+ machine:
+ resource_class: windows.medium
+ image: windows-server-2019-vs2019:stable
+ shell: powershell.exe
+
+# windows-gpu: &windows_gpu
+# machine:
+# resource_class: windows.gpu.nvidia.medium
+# image: windows-server-2019-nvidia:stable
+
+version_parameters: &version_parameters
+ parameters:
+ pytorch_version:
+ type: string
+ torchvision_version:
+ type: string
+ pytorch_index:
+ type: string
+ # use test wheels index to have access to RC wheels
+ # https://download.pytorch.org/whl/test/torch_test.html
+ default: "https://download.pytorch.org/whl/torch_stable.html"
+ python_version: # NOTE: only affect linux
+ type: string
+ default: '3.6.8'
+
+ environment:
+ PYTORCH_VERSION: << parameters.pytorch_version >>
+ TORCHVISION_VERSION: << parameters.torchvision_version >>
+ PYTORCH_INDEX: << parameters.pytorch_index >>
+ PYTHON_VERSION: << parameters.python_version>>
+ # point datasets to ~/.torch so it's cached in CI
+ DETECTRON2_DATASETS: ~/.torch/datasets
+
+# -------------------------------------------------------------------------------------
+# Re-usable commands
+# -------------------------------------------------------------------------------------
+# install_nvidia_driver: &install_nvidia_driver
+# - run:
+# name: Install nvidia driver
+# working_directory: ~/
+# command: |
+# wget -q 'https://s3.amazonaws.com/ossci-linux/nvidia_driver/NVIDIA-Linux-x86_64-430.40.run'
+# sudo /bin/bash ./NVIDIA-Linux-x86_64-430.40.run -s --no-drm
+# nvidia-smi
+
+add_ssh_keys: &add_ssh_keys
+ # https://circleci.com/docs/2.0/add-ssh-key/
+ - add_ssh_keys:
+ fingerprints:
+ - "e4:13:f2:22:d4:49:e8:e4:57:5a:ac:20:2f:3f:1f:ca"
+
+install_python: &install_python
+ - run:
+ name: Install Python
+ working_directory: ~/
+ command: |
+ # upgrade pyenv
+ cd /opt/circleci/.pyenv/plugins/python-build/../.. && git pull && cd -
+ pyenv install -s $PYTHON_VERSION
+ pyenv global $PYTHON_VERSION
+ python --version
+ which python
+ pip install --upgrade pip
+
+setup_venv: &setup_venv
+ - run:
+ name: Setup Virtual Env
+ working_directory: ~/
+ command: |
+ python -m venv ~/venv
+ echo ". ~/venv/bin/activate" >> $BASH_ENV
+ . ~/venv/bin/activate
+ python --version
+ which python
+ which pip
+ pip install --upgrade pip
+
+setup_venv_win: &setup_venv_win
+ - run:
+ name: Setup Virutal Env for Windows
+ command: |
+ pip install virtualenv
+ python -m virtualenv env
+ .\env\Scripts\activate
+ python --version
+ which python
+ which pip
+
+install_linux_dep: &install_linux_dep
+ - run:
+ name: Install Dependencies
+ command: |
+ # disable crash coredump, so unittests fail fast
+ sudo systemctl stop apport.service
+ # install from github to get latest; install iopath first since fvcore depends on it
+ pip install --progress-bar off -U 'git+https://github.com/facebookresearch/iopath'
+ pip install --progress-bar off -U 'git+https://github.com/facebookresearch/fvcore'
+ # Don't use pytest-xdist: cuda tests are unstable under multi-process workers.
+ pip install --progress-bar off ninja opencv-python-headless pytest tensorboard pycocotools
+ pip install --progress-bar off torch==$PYTORCH_VERSION -f $PYTORCH_INDEX
+ if [[ "$TORCHVISION_VERSION" == "master" ]]; then
+ pip install git+https://github.com/pytorch/vision.git
+ else
+ pip install --progress-bar off torchvision==$TORCHVISION_VERSION -f $PYTORCH_INDEX
+ fi
+
+ python -c 'import torch; print("CUDA:", torch.cuda.is_available())'
+ gcc --version
+
+install_detectron2: &install_detectron2
+ - run:
+ name: Install Detectron2
+ command: |
+ # Remove first, in case it's in the CI cache
+ pip uninstall -y detectron2
+ pip install --progress-bar off -e .[all]
+ python -m detectron2.utils.collect_env
+ ./datasets/prepare_for_tests.sh
+
+run_unittests: &run_unittests
+ - run:
+ name: Run Unit Tests
+ command: |
+ pytest -v --durations=15 tests # parallel causes some random failures
+
+# -------------------------------------------------------------------------------------
+# Jobs to run
+# -------------------------------------------------------------------------------------
+jobs:
+ linux_cpu_tests:
+ <<: *cpu
+ <<: *version_parameters
+
+ working_directory: ~/detectron2
+
+ steps:
+ - checkout
+
+ # Cache the venv directory that contains python, dependencies, and checkpoints
+ # Refresh the key when dependencies should be updated (e.g. when pytorch releases)
+ - restore_cache:
+ keys:
+ - cache-{{ arch }}-<< parameters.pytorch_version >>-{{ .Branch }}-20210827
+
+ - <<: *install_python
+ - <<: *install_linux_dep
+ - <<: *install_detectron2
+ - <<: *run_unittests
+
+ - save_cache:
+ paths:
+ - /opt/circleci/.pyenv
+ - ~/.torch
+ key: cache-{{ arch }}-<< parameters.pytorch_version >>-{{ .Branch }}-20210827
+
+
+ linux_gpu_tests:
+ <<: *gpu
+ <<: *version_parameters
+
+ working_directory: ~/detectron2
+
+ steps:
+ - checkout
+
+ - restore_cache:
+ keys:
+ - cache-{{ arch }}-<< parameters.pytorch_version >>-{{ .Branch }}-20210827
+
+ - <<: *install_python
+ - <<: *install_linux_dep
+ - <<: *install_detectron2
+ - <<: *run_unittests
+
+ - save_cache:
+ paths:
+ - /opt/circleci/.pyenv
+ - ~/.torch
+ key: cache-{{ arch }}-<< parameters.pytorch_version >>-{{ .Branch }}-20210827
+
+ windows_cpu_build:
+ <<: *windows_cpu
+ <<: *version_parameters
+ steps:
+ - <<: *add_ssh_keys
+ - checkout
+ - <<: *setup_venv_win
+
+ # Cache the env directory that contains dependencies
+ - restore_cache:
+ keys:
+ - cache-{{ arch }}-<< parameters.pytorch_version >>-{{ .Branch }}-20210404
+
+ - run:
+ name: Install Dependencies
+ command: |
+ pip install certifi --ignore-installed # required on windows to workaround some cert issue
+ pip install numpy cython # required on windows before pycocotools
+ pip install opencv-python-headless pytest-xdist pycocotools tensorboard
+ pip install -U git+https://github.com/facebookresearch/iopath
+ pip install -U git+https://github.com/facebookresearch/fvcore
+ pip install torch==$env:PYTORCH_VERSION torchvision==$env:TORCHVISION_VERSION -f $env:PYTORCH_INDEX
+
+ - save_cache:
+ paths:
+ - env
+ key: cache-{{ arch }}-<< parameters.pytorch_version >>-{{ .Branch }}-20210404
+
+ - <<: *install_detectron2
+ # TODO: unittest fails for now
+
+workflows:
+ version: 2
+ regular_test:
+ jobs:
+ - linux_cpu_tests:
+ name: linux_cpu_tests_pytorch1.10
+ pytorch_version: '1.10.0+cpu'
+ torchvision_version: '0.11.1+cpu'
+ - linux_gpu_tests:
+ name: linux_gpu_tests_pytorch1.8
+ pytorch_version: '1.8.1+cu111'
+ torchvision_version: '0.9.1+cu111'
+ - linux_gpu_tests:
+ name: linux_gpu_tests_pytorch1.9
+ pytorch_version: '1.9+cu111'
+ torchvision_version: '0.10+cu111'
+ - linux_gpu_tests:
+ name: linux_gpu_tests_pytorch1.10
+ pytorch_version: '1.10+cu111'
+ torchvision_version: '0.11.1+cu111'
+ - linux_gpu_tests:
+ name: linux_gpu_tests_pytorch1.10_python39
+ pytorch_version: '1.10+cu111'
+ torchvision_version: '0.11.1+cu111'
+ python_version: '3.9.6'
+ - windows_cpu_build:
+ pytorch_version: '1.10+cpu'
+ torchvision_version: '0.11.1+cpu'
diff --git a/models/grit_src/third_party/CenterNet2/.clang-format b/models/grit_src/third_party/CenterNet2/.clang-format
new file mode 100644
index 0000000000000000000000000000000000000000..39b1b3d603ed0cf6b7f94c9c08067f148f35613f
--- /dev/null
+++ b/models/grit_src/third_party/CenterNet2/.clang-format
@@ -0,0 +1,85 @@
+AccessModifierOffset: -1
+AlignAfterOpenBracket: AlwaysBreak
+AlignConsecutiveAssignments: false
+AlignConsecutiveDeclarations: false
+AlignEscapedNewlinesLeft: true
+AlignOperands: false
+AlignTrailingComments: false
+AllowAllParametersOfDeclarationOnNextLine: false
+AllowShortBlocksOnASingleLine: false
+AllowShortCaseLabelsOnASingleLine: false
+AllowShortFunctionsOnASingleLine: Empty
+AllowShortIfStatementsOnASingleLine: false
+AllowShortLoopsOnASingleLine: false
+AlwaysBreakAfterReturnType: None
+AlwaysBreakBeforeMultilineStrings: true
+AlwaysBreakTemplateDeclarations: true
+BinPackArguments: false
+BinPackParameters: false
+BraceWrapping:
+ AfterClass: false
+ AfterControlStatement: false
+ AfterEnum: false
+ AfterFunction: false
+ AfterNamespace: false
+ AfterObjCDeclaration: false
+ AfterStruct: false
+ AfterUnion: false
+ BeforeCatch: false
+ BeforeElse: false
+ IndentBraces: false
+BreakBeforeBinaryOperators: None
+BreakBeforeBraces: Attach
+BreakBeforeTernaryOperators: true
+BreakConstructorInitializersBeforeComma: false
+BreakAfterJavaFieldAnnotations: false
+BreakStringLiterals: false
+ColumnLimit: 80
+CommentPragmas: '^ IWYU pragma:'
+ConstructorInitializerAllOnOneLineOrOnePerLine: true
+ConstructorInitializerIndentWidth: 4
+ContinuationIndentWidth: 4
+Cpp11BracedListStyle: true
+DerivePointerAlignment: false
+DisableFormat: false
+ForEachMacros: [ FOR_EACH, FOR_EACH_R, FOR_EACH_RANGE, ]
+IncludeCategories:
+ - Regex: '^<.*\.h(pp)?>'
+ Priority: 1
+ - Regex: '^<.*'
+ Priority: 2
+ - Regex: '.*'
+ Priority: 3
+IndentCaseLabels: true
+IndentWidth: 2
+IndentWrappedFunctionNames: false
+KeepEmptyLinesAtTheStartOfBlocks: false
+MacroBlockBegin: ''
+MacroBlockEnd: ''
+MaxEmptyLinesToKeep: 1
+NamespaceIndentation: None
+ObjCBlockIndentWidth: 2
+ObjCSpaceAfterProperty: false
+ObjCSpaceBeforeProtocolList: false
+PenaltyBreakBeforeFirstCallParameter: 1
+PenaltyBreakComment: 300
+PenaltyBreakFirstLessLess: 120
+PenaltyBreakString: 1000
+PenaltyExcessCharacter: 1000000
+PenaltyReturnTypeOnItsOwnLine: 200
+PointerAlignment: Left
+ReflowComments: true
+SortIncludes: true
+SpaceAfterCStyleCast: false
+SpaceBeforeAssignmentOperators: true
+SpaceBeforeParens: ControlStatements
+SpaceInEmptyParentheses: false
+SpacesBeforeTrailingComments: 1
+SpacesInAngles: false
+SpacesInContainerLiterals: true
+SpacesInCStyleCastParentheses: false
+SpacesInParentheses: false
+SpacesInSquareBrackets: false
+Standard: Cpp11
+TabWidth: 8
+UseTab: Never
diff --git a/models/grit_src/third_party/CenterNet2/.flake8 b/models/grit_src/third_party/CenterNet2/.flake8
new file mode 100644
index 0000000000000000000000000000000000000000..ae8edda517fa7e2adc003299f18c880a823b1343
--- /dev/null
+++ b/models/grit_src/third_party/CenterNet2/.flake8
@@ -0,0 +1,15 @@
+# This is an example .flake8 config, used when developing *Black* itself.
+# Keep in sync with setup.cfg which is used for source packages.
+
+[flake8]
+ignore = W503, E203, E221, C901, C408, E741, C407, B017
+max-line-length = 100
+max-complexity = 18
+select = B,C,E,F,W,T4,B9
+exclude = build
+per-file-ignores =
+ **/__init__.py:F401,F403,E402
+ **/configs/**.py:F401,E402
+ configs/**.py:F401,E402
+ **/tests/config/**.py:F401,E402
+ tests/config/**.py:F401,E402
diff --git a/models/grit_src/third_party/CenterNet2/.github/CODE_OF_CONDUCT.md b/models/grit_src/third_party/CenterNet2/.github/CODE_OF_CONDUCT.md
new file mode 100644
index 0000000000000000000000000000000000000000..0f7ad8bfc173eac554f0b6ef7c684861e8014bbe
--- /dev/null
+++ b/models/grit_src/third_party/CenterNet2/.github/CODE_OF_CONDUCT.md
@@ -0,0 +1,5 @@
+# Code of Conduct
+
+Facebook has adopted a Code of Conduct that we expect project participants to adhere to.
+Please read the [full text](https://code.fb.com/codeofconduct/)
+so that you can understand what actions will and will not be tolerated.
diff --git a/models/grit_src/third_party/CenterNet2/.github/CONTRIBUTING.md b/models/grit_src/third_party/CenterNet2/.github/CONTRIBUTING.md
new file mode 100644
index 0000000000000000000000000000000000000000..9bab709cae689ba3b92dd52f7fbcc0c6926f4a38
--- /dev/null
+++ b/models/grit_src/third_party/CenterNet2/.github/CONTRIBUTING.md
@@ -0,0 +1,68 @@
+# Contributing to detectron2
+
+## Issues
+We use GitHub issues to track public bugs and questions.
+Please make sure to follow one of the
+[issue templates](https://github.com/facebookresearch/detectron2/issues/new/choose)
+when reporting any issues.
+
+Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## Pull Requests
+We actively welcome pull requests.
+
+However, if you're adding any significant features (e.g. > 50 lines), please
+make sure to discuss with maintainers about your motivation and proposals in an issue
+before sending a PR. This is to save your time so you don't spend time on a PR that we'll not accept.
+
+We do not always accept new features, and we take the following
+factors into consideration:
+
+1. Whether the same feature can be achieved without modifying detectron2.
+ Detectron2 is designed so that you can implement many extensions from the outside, e.g.
+ those in [projects](https://github.com/facebookresearch/detectron2/tree/master/projects).
+ * If some part of detectron2 is not extensible enough, you can also bring up a more general issue to
+ improve it. Such feature request may be useful to more users.
+2. Whether the feature is potentially useful to a large audience (e.g. an impactful detection paper, a popular dataset,
+ a significant speedup, a widely useful utility),
+ or only to a small portion of users (e.g., a less-known paper, an improvement not in the object
+ detection field, a trick that's not very popular in the community, code to handle a non-standard type of data)
+ * Adoption of additional models, datasets, new task are by default not added to detectron2 before they
+ receive significant popularity in the community.
+ We sometimes accept such features in `projects/`, or as a link in `projects/README.md`.
+3. Whether the proposed solution has a good design / interface. This can be discussed in the issue prior to PRs, or
+ in the form of a draft PR.
+4. Whether the proposed solution adds extra mental/practical overhead to users who don't
+ need such feature.
+5. Whether the proposed solution breaks existing APIs.
+
+To add a feature to an existing function/class `Func`, there are always two approaches:
+(1) add new arguments to `Func`; (2) write a new `Func_with_new_feature`.
+To meet the above criteria, we often prefer approach (2), because:
+
+1. It does not involve modifying or potentially breaking existing code.
+2. It does not add overhead to users who do not need the new feature.
+3. Adding new arguments to a function/class is not scalable w.r.t. all the possible new research ideas in the future.
+
+When sending a PR, please do:
+
+1. If a PR contains multiple orthogonal changes, split it to several PRs.
+2. If you've added code that should be tested, add tests.
+3. For PRs that need experiments (e.g. adding a new model or new methods),
+ you don't need to update model zoo, but do provide experiment results in the description of the PR.
+4. If APIs are changed, update the documentation.
+5. We use the [Google style docstrings](https://www.sphinx-doc.org/en/master/usage/extensions/napoleon.html) in python.
+6. Make sure your code lints with `./dev/linter.sh`.
+
+
+## Contributor License Agreement ("CLA")
+In order to accept your pull request, we need you to submit a CLA. You only need
+to do this once to work on any of Facebook's open source projects.
+
+Complete your CLA here: | CUDA | torch 1.10 | torch 1.9 | torch 1.8 |
|---|---|---|---|
| 11.3 | install | ||
| 11.1 | install | install | install |
| 10.2 | install | install | install |
| 10.1 | install | ||
| cpu | install | install | install |
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|
| R50-C4 | +1x | +0.551 | +0.102 | +4.8 | +35.7 | +137257644 | +model | metrics | +
| R50-DC5 | +1x | +0.380 | +0.068 | +5.0 | +37.3 | +137847829 | +model | metrics | +
| R50-FPN | +1x | +0.210 | +0.038 | +3.0 | +37.9 | +137257794 | +model | metrics | +
| R50-C4 | +3x | +0.543 | +0.104 | +4.8 | +38.4 | +137849393 | +model | metrics | +
| R50-DC5 | +3x | +0.378 | +0.070 | +5.0 | +39.0 | +137849425 | +model | metrics | +
| R50-FPN | +3x | +0.209 | +0.038 | +3.0 | +40.2 | +137849458 | +model | metrics | +
| R101-C4 | +3x | +0.619 | +0.139 | +5.9 | +41.1 | +138204752 | +model | metrics | +
| R101-DC5 | +3x | +0.452 | +0.086 | +6.1 | +40.6 | +138204841 | +model | metrics | +
| R101-FPN | +3x | +0.286 | +0.051 | +4.1 | +42.0 | +137851257 | +model | metrics | +
| X101-FPN | +3x | +0.638 | +0.098 | +6.7 | +43.0 | +139173657 | +model | metrics | +
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|
| R50 | +1x | +0.205 | +0.041 | +4.1 | +37.4 | +190397773 | +model | metrics | +
| R50 | +3x | +0.205 | +0.041 | +4.1 | +38.7 | +190397829 | +model | metrics | +
| R101 | +3x | +0.291 | +0.054 | +5.2 | +40.4 | +190397697 | +model | metrics | +
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+prop. AR |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|---|
| RPN R50-C4 | +1x | +0.130 | +0.034 | +1.5 | ++ | 51.6 | +137258005 | +model | metrics | +
| RPN R50-FPN | +1x | +0.186 | +0.032 | +2.7 | ++ | 58.0 | +137258492 | +model | metrics | +
| Fast R-CNN R50-FPN | +1x | +0.140 | +0.029 | +2.6 | +37.8 | ++ | 137635226 | +model | metrics | +
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+mask AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|---|
| R50-C4 | +1x | +0.584 | +0.110 | +5.2 | +36.8 | +32.2 | +137259246 | +model | metrics | +
| R50-DC5 | +1x | +0.471 | +0.076 | +6.5 | +38.3 | +34.2 | +137260150 | +model | metrics | +
| R50-FPN | +1x | +0.261 | +0.043 | +3.4 | +38.6 | +35.2 | +137260431 | +model | metrics | +
| R50-C4 | +3x | +0.575 | +0.111 | +5.2 | +39.8 | +34.4 | +137849525 | +model | metrics | +
| R50-DC5 | +3x | +0.470 | +0.076 | +6.5 | +40.0 | +35.9 | +137849551 | +model | metrics | +
| R50-FPN | +3x | +0.261 | +0.043 | +3.4 | +41.0 | +37.2 | +137849600 | +model | metrics | +
| R101-C4 | +3x | +0.652 | +0.145 | +6.3 | +42.6 | +36.7 | +138363239 | +model | metrics | +
| R101-DC5 | +3x | +0.545 | +0.092 | +7.6 | +41.9 | +37.3 | +138363294 | +model | metrics | +
| R101-FPN | +3x | +0.340 | +0.056 | +4.6 | +42.9 | +38.6 | +138205316 | +model | metrics | +
| X101-FPN | +3x | +0.690 | +0.103 | +7.2 | +44.3 | +39.5 | +139653917 | +model | metrics | +
| Name | +epochs | +train time (s/im) |
+inference time (s/im) |
+box AP |
+mask AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|
| R50-FPN | +100 | +0.376 | +0.069 | +44.6 | +40.3 | +42047764 | +model | metrics | +
| R50-FPN | +200 | +0.376 | +0.069 | +46.3 | +41.7 | +42047638 | +model | metrics | +
| R50-FPN | +400 | +0.376 | +0.069 | +47.4 | +42.5 | +42019571 | +model | metrics | +
| R101-FPN | +100 | +0.518 | +0.073 | +46.4 | +41.6 | +42025812 | +model | metrics | +
| R101-FPN | +200 | +0.518 | +0.073 | +48.0 | +43.1 | +42131867 | +model | metrics | +
| R101-FPN | +400 | +0.518 | +0.073 | +48.9 | +43.7 | +42073830 | +model | metrics | +
| regnetx_4gf_dds_FPN | +100 | +0.474 | +0.071 | +46.0 | +41.3 | +42047771 | +model | metrics | +
| regnetx_4gf_dds_FPN | +200 | +0.474 | +0.071 | +48.1 | +43.1 | +42132721 | +model | metrics | +
| regnetx_4gf_dds_FPN | +400 | +0.474 | +0.071 | +48.6 | +43.5 | +42025447 | +model | metrics | +
| regnety_4gf_dds_FPN | +100 | +0.487 | +0.073 | +46.1 | +41.6 | +42047784 | +model | metrics | +
| regnety_4gf_dds_FPN | +200 | +0.487 | +0.072 | +47.8 | +43.0 | +42047642 | +model | metrics | +
| regnety_4gf_dds_FPN | +400 | +0.487 | +0.072 | +48.2 | +43.3 | +42045954 | +model | metrics | +
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+kp. AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|---|
| R50-FPN | +1x | +0.315 | +0.072 | +5.0 | +53.6 | +64.0 | +137261548 | +model | metrics | +
| R50-FPN | +3x | +0.316 | +0.066 | +5.0 | +55.4 | +65.5 | +137849621 | +model | metrics | +
| R101-FPN | +3x | +0.390 | +0.076 | +6.1 | +56.4 | +66.1 | +138363331 | +model | metrics | +
| X101-FPN | +3x | +0.738 | +0.121 | +8.7 | +57.3 | +66.0 | +139686956 | +model | metrics | +
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+mask AP |
+PQ | +model id | +download | + + +
|---|---|---|---|---|---|---|---|---|---|
| R50-FPN | +1x | +0.304 | +0.053 | +4.8 | +37.6 | +34.7 | +39.4 | +139514544 | +model | metrics | +
| R50-FPN | +3x | +0.302 | +0.053 | +4.8 | +40.0 | +36.5 | +41.5 | +139514569 | +model | metrics | +
| R101-FPN | +3x | +0.392 | +0.066 | +6.0 | +42.4 | +38.5 | +43.0 | +139514519 | +model | metrics | +
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+mask AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|---|
| R50-FPN | +1x | +0.292 | +0.107 | +7.1 | +23.6 | +24.4 | +144219072 | +model | metrics | +
| R101-FPN | +1x | +0.371 | +0.114 | +7.8 | +25.6 | +25.9 | +144219035 | +model | metrics | +
| X101-FPN | +1x | +0.712 | +0.151 | +10.2 | +26.7 | +27.1 | +144219108 | +model | metrics | +
| Name | +train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+box AP50 |
+mask AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|---|
| R50-FPN, Cityscapes | +0.240 | +0.078 | +4.4 | ++ | + | 36.5 | +142423278 | +model | metrics | +
| R50-C4, VOC | +0.537 | +0.081 | +4.8 | +51.9 | +80.3 | ++ | 142202221 | +model | metrics | +
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+mask AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|---|
| Baseline R50-FPN | +1x | +0.261 | +0.043 | +3.4 | +38.6 | +35.2 | +137260431 | +model | metrics | +
| Deformable Conv | +1x | +0.342 | +0.048 | +3.5 | +41.5 | +37.5 | +138602867 | +model | metrics | +
| Cascade R-CNN | +1x | +0.317 | +0.052 | +4.0 | +42.1 | +36.4 | +138602847 | +model | metrics | +
| Baseline R50-FPN | +3x | +0.261 | +0.043 | +3.4 | +41.0 | +37.2 | +137849600 | +model | metrics | +
| Deformable Conv | +3x | +0.349 | +0.047 | +3.5 | +42.7 | +38.5 | +144998336 | +model | metrics | +
| Cascade R-CNN | +3x | +0.328 | +0.053 | +4.0 | +44.3 | +38.5 | +144998488 | +model | metrics | +
| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+mask AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|---|
| Baseline R50-FPN | +3x | +0.261 | +0.043 | +3.4 | +41.0 | +37.2 | +137849600 | +model | metrics | +
| GN | +3x | +0.309 | +0.060 | +5.6 | +42.6 | +38.6 | +138602888 | +model | metrics | +
| SyncBN | +3x | +0.345 | +0.053 | +5.5 | +41.9 | +37.8 | +169527823 | +model | metrics | +
| GN (from scratch) | +3x | +0.338 | +0.061 | +7.2 | +39.9 | +36.6 | +138602908 | +model | metrics | +
| GN (from scratch) | +9x | +N/A | +0.061 | +7.2 | +43.7 | +39.6 | +183808979 | +model | metrics | +
| SyncBN (from scratch) | +9x | +N/A | +0.055 | +7.2 | +43.6 | +39.3 | +184226666 | +model | metrics | +
| Name | +inference time (s/im) |
+train mem (GB) |
+box AP |
+mask AP |
+PQ | +model id | +download | + + +
|---|---|---|---|---|---|---|---|
| Panoptic FPN R101 | +0.098 | +11.4 | +47.4 | +41.3 | +46.1 | +139797668 | +model | metrics | +
| Mask R-CNN X152 | +0.234 | +15.1 | +50.2 | +44.0 | ++ | 18131413 | +model | metrics | +
| above + test-time aug. | ++ | + | 51.9 | +45.9 | ++ | + | + |
+
+Detectron2 is Facebook AI Research's next generation software system
+that implements state-of-the-art object detection algorithms.
+It is a ground-up rewrite of the previous version,
+[Detectron](https://github.com/facebookresearch/Detectron/),
+and it originates from [maskrcnn-benchmark](https://github.com/facebookresearch/maskrcnn-benchmark/).
+
+
+| Name | +lr sched |
+train time (s/iter) |
+inference time (s/im) |
+train mem (GB) |
+box AP |
+mask AP |
+kp. AP |
+model id | +download | + + +
|---|---|---|---|---|---|---|---|---|---|
| Faster R-CNN | +1x | +0.219 | +0.038 | +3.1 | +36.9 | ++ | + | 137781054 | +model | metrics | +
| Keypoint R-CNN | +1x | +0.313 | +0.071 | +5.0 | +53.1 | ++ | 64.2 | +137781195 | +model | metrics | +
| Mask R-CNN | +1x | +0.273 | +0.043 | +3.4 | +37.8 | +34.9 | ++ | 137781281 | +model | metrics | +