upload codes

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 OpenMMLab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

annotator/canny/__init__.py

@@ -0,0 +1,6 @@
+import cv2
+
+
+class CannyDetector:
+ def __call__(self, img, low_threshold=100, high_threshold=200):
+ return cv2.Canny(img, low_threshold, high_threshold)

annotator/cielab/__init__.py

@@ -0,0 +1,47 @@
+import os
+import sys
+sys.path.append(os.getcwd())
+sys.path.append(os.path.join(os.getcwd(), 'rayleigh'))
+
+import numpy as np
+from skimage.color import rgb2lab
+from .rayleigh import Palette
+from .rayleigh.util import histogram_colors_strict, smooth_histogram, color_hist_to_palette_image
+
+
+class CIELabDetector:
+
+ MAX_DIMENSION = 240 + 1
+
+ def __init__(self, sigma=10, num_hues=11, num_light=5, num_sat=5):
+ self.sigma = sigma
+ self.palette = Palette(num_hues=num_hues, light_range=num_light, sat_range=num_sat)
+
+ def __call__(self, img):
+ # Handle grayscale and RGBA images.
+ # TODO: Should be smarter here in the future, but for now simply remove
+ # the alpha channel if present.
+ if img.ndim == 2:
+ img = np.tile(img[:, :, np.newaxis], (1, 1, 3))
+ elif img.ndim == 4:
+ img = img[:, :, :3]
+ img = img[:,:,:3]
+
+ h, w, d = tuple(img.shape)
+ h_stride = int(h / self.MAX_DIMENSION + 1)
+ w_stride = int(w / self.MAX_DIMENSION + 1)
+ img = img[::h_stride, ::w_stride, :]
+
+ # Convert to L*a*b colors.
+ h, w, d = img.shape
+ lab_array = rgb2lab(img).reshape((h * w, d))
+
+ # compute hist
+ hist = histogram_colors_strict(lab_array, self.palette)
+ hist = smooth_histogram(hist, self.palette, self.sigma)
+ return hist
+
+ def hist_to_palette(self, hist):
+ # hist to image
+ plt = color_hist_to_palette_image(hist, self.palette)
+ return (plt * 255).astype(np.uint8)

annotator/cielab/rayleigh/__init__.py

@@ -0,0 +1,8 @@
+"""
+Rayleigh is an open-source system for quickly searching medium-sized image
+collections by multiple colors given as a palette or derived from a query image.
+"""
+
+
+from .palette import Palette
+from .util import *

annotator/cielab/rayleigh/palette.py

@@ -0,0 +1,132 @@
+"""
+Encapsulate the list of hex colors and array of Lab values representations
+of a palette (codebook) of colors.
+
+Provide methods to work with color conversion and the Palette class.
+
+Provide a parametrized method to generate a palette that covers the range
+of colors.
+"""
+
+import os
+import numpy as np
+from skimage.color import hsv2rgb, rgb2lab
+from skimage.io import imsave
+from sklearn.metrics import euclidean_distances
+
+from .util import rgb2hex
+
+
+class Palette(object):
+ """
+ Create a color palette (codebook) in the form of a 2D grid of colors,
+ as described in the parameters list below.
+ Further, the rightmost column has num_hues gradations from black to white.
+
+ Parameters
+ ----------
+ num_hues : int
+ number of colors with full lightness and saturation, in the middle
+ sat_range : int
+ number of rows above middle row that show
+ the same hues with decreasing saturation.
+ light_range : int
+ number of rows below middle row that show
+ the same hues with decreasing lightness.
+
+ Returns
+ -------
+ palette: rayleigh.Palette
+ """
+
+ def __init__(self, num_hues=8, sat_range=2, light_range=2):
+ height = 1 + sat_range + (2 * light_range - 1)
+ # generate num_hues+1 hues, but don't take the last one:
+ # hues are on a circle, and we would be oversampling the origin
+ hues = np.tile(np.linspace(0, 1, num_hues + 1)[:-1], (height, 1))
+ if num_hues == 8:
+ hues = np.tile(np.array(
+ [0., 0.10, 0.15, 0.28, 0.51, 0.58, 0.77, 0.85]), (height, 1))
+ if num_hues == 9:
+ hues = np.tile(np.array(
+ [0., 0.10, 0.15, 0.28, 0.49, 0.54, 0.60, 0.7, 0.87]), (height, 1))
+ if num_hues == 10:
+ hues = np.tile(np.array(
+ [0., 0.10, 0.15, 0.28, 0.49, 0.54, 0.60, 0.66, 0.76, 0.87]), (height, 1))
+ elif num_hues == 11:
+ hues = np.tile(np.array(
+ [0.0, 0.0833, 0.166, 0.25,
+ 0.333, 0.5, 0.56333,
+ 0.666, 0.73, 0.803,
+ 0.916]), (height, 1))
+ 
+ sats = np.hstack((
+ np.linspace(0, 1, sat_range + 2)[1:-1],
+ 1,
+ [1] * (light_range),
+ [.4] * (light_range - 1),
+ ))
+ lights = np.hstack((
+ [1] * sat_range,
+ 1,
+ np.linspace(1, 0.2, light_range + 2)[1:-1],
+ np.linspace(1, 0.2, light_range + 2)[1:-2],
+ ))
+
+ sats = np.tile(np.atleast_2d(sats).T, (1, num_hues))
+ lights = np.tile(np.atleast_2d(lights).T, (1, num_hues))
+ colors = hsv2rgb(np.dstack((hues, sats, lights)))
+ grays = np.tile(
+ np.linspace(1, 0, height)[:, np.newaxis, np.newaxis], (1, 1, 3))
+
+ self.rgb_image = np.hstack((colors, grays))
+
+ # Make a nice histogram ordering of the hues and grays
+ h, w, d = colors.shape
+ color_array = colors.T.reshape((d, w * h)).T
+ h, w, d = grays.shape
+ gray_array = grays.T.reshape((d, w * h)).T
+ 
+ self.rgb_array = np.vstack((color_array, gray_array))
+ self.lab_array = rgb2lab(self.rgb_array[None, :, :]).squeeze()
+ self.hex_list = [rgb2hex(row) for row in self.rgb_array]
+ #assert(np.all(self.rgb_array == self.rgb_array[None, :, :].squeeze()))
+
+ self.distances = euclidean_distances(self.lab_array, squared=True)
+
+ def output(self, dirname, html=False):
+ """
+ Output an image of the palette, josn list of the hex
+ colors, and an HTML color picker for it.
+
+ Parameters
+ ----------
+ dirname : string
+ directory for the files to be output
+ """
+ def get_palette_html():
+ """
+ Return HTML for a color picker using the given palette.
+ """
+ html = """
+ <style>
+ span {
+ width: 20px;
+ height: 20px;
+ margin: 2px;
+ padding: 0px;
+ display: inline-block;
+ }
+ </style>
+ """
+ for row in self.rgb_image:
+ for rgb_color in row:
+ s = '<a id="{0}"><span style="background-color: {0}" /></a>\n'
+ html += s.format(rgb2hex(rgb_color))
+ html += "<br />\n"
+ return html
+
+ imsave(os.path.join(dirname, 'palette.png'), (self.rgb_image*255).astype(np.uint8))
+ if html:
+ with open(os.path.join(dirname, 'palette.html'), 'w') as f:
+ f.write(get_palette_html())

annotator/cielab/rayleigh/util.py

@@ -0,0 +1,270 @@
+import os
+import numpy as np
+import tempfile
+import matplotlib.pyplot as plt
+from sklearn.metrics import euclidean_distances
+from skimage.io import imsave
+
+
+def rgb2hex(rgb_number):
+ """
+ Args:
+ - rgb_number (sequence of float)
+
+ Returns:
+ - hex_number (string)
+ """
+ return '#%02x%02x%02x' % tuple([int(np.round(val * 255)) for val in rgb_number])
+
+
+def hex2rgb(hexcolor_str):
+ """
+ Args:
+ - hexcolor_str (string): e.g. '#ffffff' or '33cc00'
+
+ Returns:
+ - rgb_color (sequence of floats): e.g. (0.2, 0.3, 0)
+ """
+ color = hexcolor_str.strip('#')
+ rgb = lambda x: round(int(x, 16) / 255., 5)
+ return (rgb(color[:2]), rgb(color[2:4]), rgb(color[4:6]))
+
+
+def color_hist_to_palette_image(color_hist, palette, percentile=90,
+ width=200, height=50, filename=None):
+ """
+ Output the main colors in the histogram to a "palette image."
+
+ Parameters
+ ----------
+ color_hist : (K,) ndarray
+ palette : rayleigh.Palette
+ percentile : int, optional:
+ Output only colors above this percentile of prevalence in the histogram.
+ filename : string, optional:
+ If given, save the resulting image to file.
+
+ Returns
+ -------
+ rgb_image : ndarray
+ """
+ ind = np.argsort(-color_hist)
+ ind = ind[color_hist[ind] > np.percentile(color_hist, percentile)]
+ hex_list = np.take(palette.hex_list, ind)
+ values = color_hist[ind]
+ rgb_image = palette_query_to_rgb_image(dict(zip(hex_list, values)))
+ if filename:
+ imsave(filename, rgb_image)
+ return rgb_image
+
+
+def palette_query_to_rgb_image(palette_query, width=200, height=50):
+ """
+ Convert a list of hex colors and their values to an RGB image of given
+ width and height.
+
+ Args:
+ - palette_query (dict):
+ a dictionary of hex colors to unnormalized values,
+ e.g. {'#ffffff': 20, '#33cc00': 0.4}.
+ """
+ hex_list, values = zip(*palette_query.items())
+ values = np.array(values)
+ values /= values.sum()
+ nums = np.array(values * width, dtype=int)
+ rgb_arrays = (np.tile(np.array(hex2rgb(x)), (num, 1))
+ for x, num in zip(hex_list, nums))
+ rgb_array = np.vstack(list(rgb_arrays))
+ rgb_image = rgb_array[np.newaxis, :, :]
+ rgb_image = np.tile(rgb_image, (height, 1, 1))
+ return rgb_image
+
+
+def plot_histogram(color_hist, palette, plot_filename=None):
+ """
+ Return Figure containing the color palette histogram.
+
+ Args:
+ - color_hist (K, ndarray)
+
+ - palette (Palette)
+
+ - plot_filename (string) [default=None]:
+ Save histogram to this file, if given.
+
+ Returns:
+ - fig (Figure)
+ """
+ fig = plt.figure(figsize=(5, 3), dpi=150)
+ ax = fig.add_subplot(111)
+ ax.bar(
+ range(len(color_hist)), color_hist,
+ color=palette.hex_list, edgecolor='black')
+ ax.set_ylim((0, 0.3))
+ ax.xaxis.set_ticks([])
+ ax.set_xlim((0, len(palette.hex_list)))
+ if plot_filename:
+ fig.savefig(plot_filename, dpi=150, facecolor='none')
+ return fig
+
+
+def output_histogram_base64(color_hist, palette):
+ """
+ Return base64-encoded image containing the color palette histogram.
+
+ Args:
+ - color_hist (K, ndarray)
+
+ - palette (Palette)
+
+ Returns:
+ - data_uri (base64 encoded string)
+ """
+ _, tfname = tempfile.mkstemp('.png')
+ plot_histogram(color_hist, palette, tfname)
+ data_uri = open(tfname, 'rb').read().encode('base64').replace('\n', '')
+ os.remove(tfname)
+ return data_uri
+
+
+def histogram_colors_strict(lab_array, palette, plot_filename=None):
+ """
+ Return a palette histogram of colors in the image.
+
+ Parameters
+ ----------
+ lab_array : (N,3) ndarray
+ The L*a*b color of each of N pixels.
+ palette : rayleigh.Palette
+ Containing K colors.
+ plot_filename : string, optional
+ If given, save histogram to this filename.
+
+ Returns
+ -------
+ color_hist : (K,) ndarray
+ """
+ # This is the fastest way that I've found.
+ # >>> %%timeit -n 200 from sklearn.metrics import euclidean_distances
+ # >>> euclidean_distances(palette, lab_array, squared=True)
+ dist = euclidean_distances(palette.lab_array, lab_array, squared=True).T
+ min_ind = np.argmin(dist, axis=1)
+ num_colors = palette.lab_array.shape[0]
+ num_pixels = lab_array.shape[0]
+ color_hist = 1. * np.bincount(min_ind, minlength=num_colors) / num_pixels
+ if plot_filename is not None:
+ plot_histogram(color_hist, palette, plot_filename)
+ return color_hist
+
+
+def histogram_colors_smoothed(lab_array, palette, sigma=10,
+ plot_filename=None, direct=True):
+ """
+ Returns a palette histogram of colors in the image, smoothed with
+ a Gaussian. Can smooth directly per-pixel, or after computing a strict
+ histogram.
+
+ Parameters
+ ----------
+ lab_array : (N,3) ndarray
+ The L*a*b color of each of N pixels.
+ palette : rayleigh.Palette
+ Containing K colors.
+ sigma : float
+ Variance of the smoothing Gaussian.
+ direct : bool, optional
+ If True, constructs a smoothed histogram directly from pixels.
+ If False, constructs a nearest-color histogram and then smoothes it.
+
+ Returns
+ -------
+ color_hist : (K,) ndarray
+ """
+ if direct:
+ color_hist_smooth = histogram_colors_with_smoothing(
+ lab_array, palette, sigma)
+ else:
+ color_hist_strict = histogram_colors_strict(lab_array, palette)
+ color_hist_smooth = smooth_histogram(color_hist_strict, palette, sigma)
+ if plot_filename is not None:
+ plot_histogram(color_hist_smooth, palette, plot_filename)
+ return color_hist_smooth
+
+
+def smooth_histogram(color_hist, palette, sigma=10):
+ """
+ Smooth the given palette histogram with a Gaussian of variance sigma.
+
+ Parameters
+ ----------
+ color_hist : (K,) ndarray
+ palette : rayleigh.Palette
+ containing K colors.
+
+ Returns
+ -------
+ color_hist_smooth : (K,) ndarray
+ """
+ n = 2. * sigma ** 2
+ weights = np.exp(-palette.distances / n)
+ norm_weights = weights / weights.sum(1)[:, np.newaxis]
+ color_hist_smooth = (norm_weights * color_hist).sum(1)
+ color_hist_smooth[color_hist_smooth < 1e-5] = 0
+ return color_hist_smooth
+
+
+def histogram_colors_with_smoothing(lab_array, palette, sigma=10):
+ """
+ Assign colors in the image to nearby colors in the palette, weighted by
+ distance in Lab color space.
+
+ Parameters
+ ----------
+ lab_array (N,3) ndarray:
+ N is the number of data points, columns are L, a, b values.
+ palette : rayleigh.Palette
+ containing K colors.
+ sigma : float
+ (0,1] value to control the steepness of exponential falloff.
+ To see the effect:
+
+ >>> from pylab import *
+ >>> ds = linspace(0,5000) # squared distance
+ >>> sigma=10; plot(ds, exp(-ds/(2*sigma**2)), label='$\sigma=%.1f$'%sigma)
+ >>> sigma=20; plot(ds, exp(-ds/(2*sigma**2)), label='$\sigma=%.1f$'%sigma)
+ >>> sigma=40; plot(ds, exp(-ds/(2*sigma**2)), label='$\sigma=%.1f$'%sigma)
+ >>> ylim([0,1]); legend();
+ >>> xlabel('Squared distance'); ylabel('Weight');
+ >>> title('Exponential smoothing')
+ >>> #plt.savefig('exponential_smoothing.png', dpi=300)
+
+ sigma=20 seems reasonable: hits 0 around squared distance of 4000.
+
+ Returns:
+ color_hist : (K,) ndarray
+ the normalized, smooth histogram of colors.
+ """
+ dist = euclidean_distances(palette.lab_array, lab_array, squared=True).T
+ n = 2. * sigma ** 2
+ weights = np.exp(-dist / n)
+ 
+ # normalize by sum: if a color is equally well represented by several colors
+ # it should not contribute much to the overall histogram
+ normalizing = weights.sum(1)
+ normalizing[normalizing == 0] = 1e16
+ normalized_weights = weights / normalizing[:, np.newaxis]
+
+ color_hist = normalized_weights.sum(0)
+ color_hist /= lab_array.shape[0]
+ color_hist[color_hist < 1e-5] = 0
+ return color_hist
+
+
+def makedirs(dirname):
+ "Does what mkdir -p does, and returns dirname."
+ if not os.path.exists(dirname):
+ try:
+ os.makedirs(dirname)
+ except:
+ print("Exception on os.makedirs")
+ return dirname

annotator/content/__init__.py

@@ -0,0 +1,23 @@
+import cv2
+import numpy as np
+from PIL import Image
+
+import torch
+from transformers import AutoProcessor, CLIPModel
+
+from annotator.util import annotator_ckpts_path
+
+
+class ContentDetector:
+ def __init__(self, model_name="openai/clip-vit-large-patch14"):
+
+ self.model = CLIPModel.from_pretrained(model_name, cache_dir=annotator_ckpts_path).cuda().eval()
+ self.processor = AutoProcessor.from_pretrained(model_name, cache_dir=annotator_ckpts_path)
+
+ def __call__(self, img):
+ with torch.no_grad():
+ img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+ inputs = self.processor(images=[img], return_tensors="pt").to('cuda')
+ image_features = self.model.get_image_features(**inputs)
+ content_emb = image_features[0].detach().cpu().numpy()
+ return content_emb

annotator/entityseg/__init__.py

@@ -0,0 +1,93 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from: https://github.com/facebookresearch/detectron2/blob/master/demo/demo.py
+import argparse
+import glob
+import multiprocessing as mp
+import os
+import sys
+sys.path.insert(1, os.getcwd())
+
+import tempfile
+import time
+import warnings
+
+import cv2
+import numpy as np
+import tqdm
+import torch
+
+from detectron2.config import get_cfg
+from detectron2.data.detection_utils import read_image
+from detectron2.projects.deeplab import add_deeplab_config
+from detectron2.utils.logger import setup_logger
+
+from mask2former import add_maskformer2_config
+from predictor import VisualizationDemo
+
+from annotator.util import annotator_ckpts_path
+
+
+model_url = "https://huggingface.co/datasets/qqlu1992/Adobe_EntitySeg/resolve/main/CropFormer_model/Entity_Segmentation/CropFormer_hornet_3x.pth"
+
+
+def make_colors():
+ from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES
+ colors = []
+ for cate in COCO_CATEGORIES:
+ colors.append(cate["color"])
+ return colors
+
+
+class EntitysegDetector:
+
+ def __init__(self, confidence_threshold=0.5):
+ cfg = get_cfg()
+ add_deeplab_config(cfg)
+ add_maskformer2_config(cfg)
+
+ workdir = os.getcwd()
+ config_file = f"{workdir}/annotator/entityseg/configs/cropformer_hornet_3x.yaml"
+ model_path = f'{annotator_ckpts_path}/CropFormer_hornet_3x_03823a.pth'
+ # Authentication required
+ # if not os.path.exists(model_path):
+ # from basicsr.utils.download_util import load_file_from_url
+ # load_file_from_url(model_url, model_dir=annotator_ckpts_path)
+
+ cfg.merge_from_file(config_file)
+ opts = ['MODEL.WEIGHTS', model_path]
+ cfg.merge_from_list(opts)
+ cfg.freeze()
+
+ self.confidence_threshold = confidence_threshold
+
+ self.colors = make_colors()
+ self.demo = VisualizationDemo(cfg)
+
+
+ def __call__(self, image): 
+ predictions = self.demo.run_on_image(image)
+ ##### color_mask
+ pred_masks = predictions["instances"].pred_masks
+ pred_scores = predictions["instances"].scores
+ 
+ # select by confidence threshold
+ selected_indexes = (pred_scores >= self.confidence_threshold)
+ selected_scores = pred_scores[selected_indexes]
+ selected_masks = pred_masks[selected_indexes]
+ _, m_H, m_W = selected_masks.shape
+ mask_id = np.zeros((m_H, m_W), dtype=np.uint8)
+
+ # rank
+ selected_scores, ranks = torch.sort(selected_scores)
+ ranks = ranks + 1
+ for index in ranks:
+ mask_id[(selected_masks[index-1]==1).cpu().numpy()] = int(index)
+ unique_mask_id = np.unique(mask_id)
+
+ color_mask = np.zeros(image.shape, dtype=np.uint8)
+ for count in unique_mask_id:
+ if count == 0:
+ continue
+ color_mask[mask_id==count] = self.colors[count % len(self.colors)]
+ 
+ return color_mask

annotator/entityseg/configs/Base-Mask2Former.yaml

@@ -0,0 +1,49 @@
+ENTITY:
+ ENABLE: True
+MODEL:
+ BACKBONE:
+ FREEZE_AT: 0
+ NAME: "build_resnet_backbone"
+ WEIGHTS: "R-50.pkl"
+ PIXEL_MEAN: [123.675, 116.280, 103.530]
+ PIXEL_STD: [58.395, 57.120, 57.375]
+ RESNETS:
+ DEPTH: 50
+ STEM_TYPE: "basic" # not used
+ STEM_OUT_CHANNELS: 64
+ STRIDE_IN_1X1: False
+ OUT_FEATURES: ["res2", "res3", "res4", "res5"]
+ # NORM: "SyncBN"
+ RES5_MULTI_GRID: [1, 1, 1] # not used
+DATASETS:
+ TRAIN: ("entityv2_entity_train_01",)
+ TEST: ("entityv2_entity_val_01",)
+SOLVER:
+ STEPS: (30525, 33138)
+ MAX_ITER: 34375
+ IMS_PER_BATCH: 16
+ BASE_LR: 0.0001
+ WARMUP_FACTOR: 1.0
+ WARMUP_ITERS: 0
+ WEIGHT_DECAY: 0.05
+ OPTIMIZER: "ADAMW"
+ LR_SCHEDULER_NAME: "WarmupPolyLR"
+ BACKBONE_MULTIPLIER: 0.1
+ CLIP_GRADIENTS:
+ ENABLED: True
+ CLIP_TYPE: "full_model"
+ CLIP_VALUE: 0.01
+ NORM_TYPE: 2.0
+ AMP:
+ ENABLED: True
+INPUT:
+ MASK_FORMAT: "bitmask"
+ FORMAT: "RGB"
+ MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800)
+ DATASET_MAPPER_NAME: "entity_crop"
+TEST:
+ EVAL_PERIOD: 400000
+DATALOADER:
+ FILTER_EMPTY_ANNOTATIONS: True
+ NUM_WORKERS: 32
+VERSION: 2

annotator/entityseg/configs/cropformer_hornet_3x.yaml

@@ -0,0 +1,70 @@
+_BASE_: Base-Mask2Former.yaml
+DATALOADER:
+ NUM_WORKERS: 32
+DATASETS:
+ TRAIN: ("entityv2_entity_train_01","entityv2_entity_train_02","entityv2_entity_train_03",)
+ TEST: ("entityv2_entity_val_all",)
+ # TEST: ("entityv2_entity_val_all_lr",)
+SOLVER:
+ # STEPS: (91575, 99414)
+ # MAX_ITER: 103125
+ IMS_PER_BATCH: 8
+ STEPS: (183150, 198828)
+ MAX_ITER: 206250
+MODEL:
+ BACKBONE:
+ NAME: "D2HorNet"
+ PIXEL_MEAN: [123.675, 116.28, 103.53]
+ PIXEL_STD: [58.395, 57.120, 57.375]
+ SWIN:
+ EMBED_DIM: 192
+ DEPTHS: [2, 2, 18, 2]
+ NUM_HEADS: [6, 12, 24, 48]
+ WINDOW_SIZE: 7
+ APE: False
+ DROP_PATH_RATE: 0.3
+ PATCH_NORM: True
+ PRETRAIN_IMG_SIZE: 384
+ WEIGHTS: "hornet_l_pretrained.pth"
+ META_ARCHITECTURE: "CropFormer"
+ SEM_SEG_HEAD:
+ NAME: "MaskFormerHead"
+ IGNORE_VALUE: 255
+ NUM_CLASSES: 1
+ LOSS_WEIGHT: 1.0
+ CONVS_DIM: 256
+ MASK_DIM: 256
+ NORM: "GN"
+ # pixel decoder
+ PIXEL_DECODER_NAME: "MSDeformAttnPixelDecoder"
+ IN_FEATURES: ["res2", "res3", "res4", "res5"]
+ DEFORMABLE_TRANSFORMER_ENCODER_IN_FEATURES: ["res3", "res4", "res5"]
+ COMMON_STRIDE: 4
+ TRANSFORMER_ENC_LAYERS: 6
+ MASK_FORMER:
+ TRANSFORMER_DECODER_NAME: "CropSharedMultiScaleMaskedTransformerDecoder"
+ TRANSFORMER_IN_FEATURE: "multi_scale_pixel_decoder"
+ DEEP_SUPERVISION: True
+ NO_OBJECT_WEIGHT: 0.1
+ CLASS_WEIGHT: 2.0
+ MASK_WEIGHT: 5.0
+ DICE_WEIGHT: 5.0
+ HIDDEN_DIM: 256
+ NUM_OBJECT_QUERIES: 200
+ NHEADS: 8
+ DROPOUT: 0.0
+ DIM_FEEDFORWARD: 2048
+ ENC_LAYERS: 0
+ PRE_NORM: False
+ ENFORCE_INPUT_PROJ: False
+ SIZE_DIVISIBILITY: 32
+ DEC_LAYERS: 10 # 9 decoder layers, add one for the loss on learnable query
+ TRAIN_NUM_POINTS: 12544
+ OVERSAMPLE_RATIO: 3.0
+ IMPORTANCE_SAMPLE_RATIO: 0.75
+ TEST:
+ SEMANTIC_ON: False
+ INSTANCE_ON: True
+ PANOPTIC_ON: False
+ OVERLAP_THRESHOLD: 0.8
+ OBJECT_MASK_THRESHOLD: 0.8

annotator/entityseg/mask2former/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from . import data # register all new datasets
+from . import modeling
+
+# config
+from .config import add_maskformer2_config
+
+# models
+from .maskformer_model import MaskFormer
+from .cropformer_model import CropFormer
+from .test_time_augmentation import SemanticSegmentorWithTTA

annotator/entityseg/mask2former/config.py

@@ -0,0 +1,139 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+from detectron2.config import CfgNode as CN
+
+
+def add_maskformer2_config(cfg):
+ """
+ Add config for MASK_FORMER.
+ """
+ # NOTE: configs from original maskformer
+ # data config
+ # select the dataset mapper
+ cfg.INPUT.DATASET_MAPPER_NAME = "mask_former_semantic"
+ # Color augmentation
+ cfg.INPUT.COLOR_AUG_SSD = False
+ # We retry random cropping until no single category in semantic segmentation GT occupies more
+ # than `SINGLE_CATEGORY_MAX_AREA` part of the crop.
+ cfg.INPUT.CROP.SINGLE_CATEGORY_MAX_AREA = 1.0
+ # Pad image and segmentation GT in dataset mapper.
+ cfg.INPUT.SIZE_DIVISIBILITY = -1
+
+ # solver config
+ # weight decay on embedding
+ cfg.SOLVER.WEIGHT_DECAY_EMBED = 0.0
+ # optimizer
+ cfg.SOLVER.OPTIMIZER = "ADAMW"
+ cfg.SOLVER.BACKBONE_MULTIPLIER = 0.1
+
+ # mask_former model config
+ cfg.MODEL.MASK_FORMER = CN()
+
+ # loss
+ cfg.MODEL.MASK_FORMER.DEEP_SUPERVISION = True
+ cfg.MODEL.MASK_FORMER.NO_OBJECT_WEIGHT = 0.1
+ cfg.MODEL.MASK_FORMER.CLASS_WEIGHT = 1.0
+ cfg.MODEL.MASK_FORMER.DICE_WEIGHT = 1.0
+ cfg.MODEL.MASK_FORMER.MASK_WEIGHT = 20.0
+
+ # transformer config
+ cfg.MODEL.MASK_FORMER.NHEADS = 8
+ cfg.MODEL.MASK_FORMER.DROPOUT = 0.1
+ cfg.MODEL.MASK_FORMER.DIM_FEEDFORWARD = 2048
+ cfg.MODEL.MASK_FORMER.ENC_LAYERS = 0
+ cfg.MODEL.MASK_FORMER.DEC_LAYERS = 6
+ cfg.MODEL.MASK_FORMER.PRE_NORM = False
+
+ cfg.MODEL.MASK_FORMER.HIDDEN_DIM = 256
+ cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES = 100
+
+ cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE = "res5"
+ cfg.MODEL.MASK_FORMER.ENFORCE_INPUT_PROJ = False
+
+ # mask_former inference config
+ cfg.MODEL.MASK_FORMER.TEST = CN()
+ cfg.MODEL.MASK_FORMER.TEST.SEMANTIC_ON = True
+ cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON = False
+ cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON = False
+ cfg.MODEL.MASK_FORMER.TEST.OBJECT_MASK_THRESHOLD = 0.0
+ cfg.MODEL.MASK_FORMER.TEST.OVERLAP_THRESHOLD = 0.0
+ cfg.MODEL.MASK_FORMER.TEST.SEM_SEG_POSTPROCESSING_BEFORE_INFERENCE = False
+
+ # Sometimes `backbone.size_divisibility` is set to 0 for some backbone (e.g. ResNet)
+ # you can use this config to override
+ cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY = 32
+
+ # pixel decoder config
+ cfg.MODEL.SEM_SEG_HEAD.MASK_DIM = 256
+ # adding transformer in pixel decoder
+ cfg.MODEL.SEM_SEG_HEAD.TRANSFORMER_ENC_LAYERS = 0
+ # pixel decoder
+ cfg.MODEL.SEM_SEG_HEAD.PIXEL_DECODER_NAME = "BasePixelDecoder"
+
+ # swin transformer backbone
+ cfg.MODEL.SWIN = CN()
+ cfg.MODEL.SWIN.PRETRAIN_IMG_SIZE = 224
+ cfg.MODEL.SWIN.PATCH_SIZE = 4
+ cfg.MODEL.SWIN.EMBED_DIM = 96
+ cfg.MODEL.SWIN.DEPTHS = [2, 2, 6, 2]
+ cfg.MODEL.SWIN.NUM_HEADS = [3, 6, 12, 24]
+ cfg.MODEL.SWIN.WINDOW_SIZE = 7
+ cfg.MODEL.SWIN.MLP_RATIO = 4.0
+ cfg.MODEL.SWIN.QKV_BIAS = True
+ cfg.MODEL.SWIN.QK_SCALE = None
+ cfg.MODEL.SWIN.DROP_RATE = 0.0
+ cfg.MODEL.SWIN.ATTN_DROP_RATE = 0.0
+ cfg.MODEL.SWIN.DROP_PATH_RATE = 0.3
+ cfg.MODEL.SWIN.APE = False
+ cfg.MODEL.SWIN.PATCH_NORM = True
+ cfg.MODEL.SWIN.OUT_FEATURES = ["res2", "res3", "res4", "res5"]
+ cfg.MODEL.SWIN.USE_CHECKPOINT = False
+
+ # NOTE: maskformer2 extra configs
+ # transformer module
+ cfg.MODEL.MASK_FORMER.TRANSFORMER_DECODER_NAME = "MultiScaleMaskedTransformerDecoder"
+
+ # LSJ aug
+ cfg.INPUT.IMAGE_SIZE = 1024
+ cfg.INPUT.MIN_SCALE = 0.1
+ cfg.INPUT.MAX_SCALE = 2.0
+
+ # MSDeformAttn encoder configs
+ cfg.MODEL.SEM_SEG_HEAD.DEFORMABLE_TRANSFORMER_ENCODER_IN_FEATURES = ["res3", "res4", "res5"]
+ cfg.MODEL.SEM_SEG_HEAD.DEFORMABLE_TRANSFORMER_ENCODER_N_POINTS = 4
+ cfg.MODEL.SEM_SEG_HEAD.DEFORMABLE_TRANSFORMER_ENCODER_N_HEADS = 8
+
+ # point loss configs
+ # Number of points sampled during training for a mask point head.
+ cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS = 112 * 112
+ # Oversampling parameter for PointRend point sampling during training. Parameter `k` in the
+ # original paper.
+ cfg.MODEL.MASK_FORMER.OVERSAMPLE_RATIO = 3.0
+ # Importance sampling parameter for PointRend point sampling during training. Parametr `beta` in
+ # the original paper.
+ cfg.MODEL.MASK_FORMER.IMPORTANCE_SAMPLE_RATIO = 0.75
+
+ ## For Entity
+ cfg.ENTITY = CN()
+ cfg.ENTITY.ENABLE = False
+ cfg.ENTITY.CROP_AREA_RATIO = 0.7
+ cfg.ENTITY.CROP_STRIDE_RATIO = 0.6
+ cfg.ENTITY.CROP_SAMPLE_NUM_TRAIN = 1
+ cfg.ENTITY.CROP_SAMPLE_NUM_TEST = 4
+
+ ## fuse frame embeddings to batch embedding
+ cfg.ENTITY.FUSE_NUM_LAYERS = 1
+ cfg.ENTITY.FUSE_ENC_HIDDIEN_DIM = 256
+ cfg.ENTITY.FUSE_ENC_NHEADS = 8
+ cfg.ENTITY.FUSE_ENC_PRE_NORM = False
+ cfg.ENTITY.FUSE_ENC_DIM_FEEDFORWARD = 2048
+ cfg.ENTITY.FUSE_ENC_LAST_LAYERS = 1
+ cfg.ENTITY.FUSE_DEC_NUM_LAYERS = 3
+
+ ## Hornet backbone
+ cfg.MODEL.HORNET = CN()
+ cfg.MODEL.HORNET.DEPTHS = [2, 3, 18, 2]
+ cfg.MODEL.HORNET.BASE_DIM = 192
+ cfg.MODEL.HORNET.GCONV = ['partial(gnconv, order=2, s=1/3)', 'partial(gnconv, order=3, s=1/3)', 'partial(gnconv, order=4, s=1/3, h=24, w=13, gflayer=GlobalLocalFilter)', 'partial(gnconv, order=5, s=1/3, h=12, w=7, gflayer=GlobalLocalFilter)']
+ cfg.MODEL.HORNET.DROP_PATH_RATE=0.6
+ cfg.MODEL.HORNET.OUT_FEATURES = ["res2", "res3", "res4", "res5"]

annotator/entityseg/mask2former/cropformer_model.py

@@ -0,0 +1,678 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from typing import Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+import pdb
+import numpy as np
+import cv2
+import os
+
+from detectron2.config import configurable
+from detectron2.data import MetadataCatalog
+from detectron2.modeling import META_ARCH_REGISTRY, build_backbone, build_sem_seg_head
+from detectron2.modeling.backbone import Backbone
+from detectron2.modeling.postprocessing import sem_seg_postprocess
+from detectron2.structures import Boxes, ImageList, Instances, BitMasks
+from detectron2.utils.memory import retry_if_cuda_oom
+from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES
+
+from .modeling.criterion import SetCriterion
+from .modeling.matcher import HungarianMatcher
+from .modeling.criterion_view import ViewSetCriterion
+from .modeling.matcher_view import ViewHungarianMatcher
+import pdb
+import copy
+
+@META_ARCH_REGISTRY.register()
+class CropFormer(nn.Module):
+ """
+ Main class for mask classification semantic segmentation architectures.
+ """
+ @configurable
+ def __init__(
+ self,
+ *,
+ cfg,
+ backbone: Backbone,
+ sem_seg_head: nn.Module,
+ criterion_2d: nn.Module,
+ criterion_3d: nn.Module,
+ num_queries: int,
+ object_mask_threshold: float,
+ overlap_threshold: float,
+ metadata,
+ size_divisibility: int,
+ sem_seg_postprocess_before_inference: bool,
+ pixel_mean: Tuple[float],
+ pixel_std: Tuple[float],
+ # inference
+ semantic_on: bool,
+ panoptic_on: bool,
+ instance_on: bool,
+ test_topk_per_image: int,
+ ):
+ """
+ Args:
+ backbone: a backbone module, must follow detectron2's backbone interface
+ sem_seg_head: a module that predicts semantic segmentation from backbone features
+ criterion: a module that defines the loss
+ num_queries: int, number of queries
+ object_mask_threshold: float, threshold to filter query based on classification score
+ for panoptic segmentation inference
+ overlap_threshold: overlap threshold used in general inference for panoptic segmentation
+ metadata: dataset meta, get `thing` and `stuff` category names for panoptic
+ segmentation inference
+ size_divisibility: Some backbones require the input height and width to be divisible by a
+ specific integer. We can use this to override such requirement.
+ sem_seg_postprocess_before_inference: whether to resize the prediction back
+ to original input size before semantic segmentation inference or after.
+ For high-resolution dataset like Mapillary, resizing predictions before
+ inference will cause OOM error.
+ pixel_mean, pixel_std: list or tuple with #channels element, representing
+ the per-channel mean and std to be used to normalize the input image
+ semantic_on: bool, whether to output semantic segmentation prediction
+ instance_on: bool, whether to output instance segmentation prediction
+ panoptic_on: bool, whether to output panoptic segmentation prediction
+ test_topk_per_image: int, instance segmentation parameter, keep topk instances per image
+ """
+ super().__init__()
+ self.cfg = cfg
+ self.backbone = backbone
+ self.sem_seg_head = sem_seg_head
+ self.criterion_2d = criterion_2d
+ self.criterion_3d = criterion_3d
+ ## colors
+ self.colors = [info["color"] for info in COCO_CATEGORIES]
+
+ self.num_queries = num_queries
+ self.overlap_threshold = overlap_threshold
+ self.object_mask_threshold = object_mask_threshold
+ self.metadata = metadata
+ if size_divisibility < 0:
+ # use backbone size_divisibility if not set
+ size_divisibility = self.backbone.size_divisibility
+ self.size_divisibility = size_divisibility
+ self.sem_seg_postprocess_before_inference = sem_seg_postprocess_before_inference
+ self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False)
+ self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)
+
+ ## colors
+ self.colors = [info["color"] for info in COCO_CATEGORIES]
+
+ # additional args
+ self.semantic_on = semantic_on
+ self.instance_on = instance_on
+ self.panoptic_on = panoptic_on
+ self.test_topk_per_image = test_topk_per_image
+
+ if not self.semantic_on:
+ assert self.sem_seg_postprocess_before_inference
+
+ @classmethod
+ def from_config(cls, cfg):
+ backbone = build_backbone(cfg)
+ sem_seg_head = build_sem_seg_head(cfg, backbone.output_shape())
+
+ # Loss parameters:
+ deep_supervision = cfg.MODEL.MASK_FORMER.DEEP_SUPERVISION
+ no_object_weight = cfg.MODEL.MASK_FORMER.NO_OBJECT_WEIGHT
+
+ # loss weights
+ class_weight = cfg.MODEL.MASK_FORMER.CLASS_WEIGHT
+ dice_weight = cfg.MODEL.MASK_FORMER.DICE_WEIGHT
+ mask_weight = cfg.MODEL.MASK_FORMER.MASK_WEIGHT
+
+ # building criterion
+ matcher_2d = HungarianMatcher(
+ cost_class=class_weight,
+ cost_mask=mask_weight,
+ cost_dice=dice_weight,
+ num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS,
+ )
+
+ matcher_3d = ViewHungarianMatcher(
+ cost_class=class_weight,
+ cost_mask=mask_weight,
+ cost_dice=dice_weight,
+ num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS,
+ )
+
+ weight_dict = {"loss_ce": class_weight, "loss_mask": mask_weight, "loss_dice": dice_weight}
+
+ if deep_supervision:
+ dec_layers = cfg.MODEL.MASK_FORMER.DEC_LAYERS
+ aux_weight_dict = {}
+ for i in range(dec_layers - 1):
+ aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
+ weight_dict.update(aux_weight_dict)
+
+ losses = ["labels", "masks"]
+
+ criterion_2d = SetCriterion(
+ sem_seg_head.num_classes,
+ matcher=matcher_2d,
+ weight_dict=weight_dict,
+ eos_coef=no_object_weight,
+ losses=losses,
+ num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS,
+ oversample_ratio=cfg.MODEL.MASK_FORMER.OVERSAMPLE_RATIO,
+ importance_sample_ratio=cfg.MODEL.MASK_FORMER.IMPORTANCE_SAMPLE_RATIO,
+ )
+
+ criterion_3d = ViewSetCriterion(
+ sem_seg_head.num_classes,
+ matcher=matcher_3d,
+ weight_dict=weight_dict,
+ eos_coef=no_object_weight,
+ losses=losses,
+ num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS,
+ oversample_ratio=cfg.MODEL.MASK_FORMER.OVERSAMPLE_RATIO,
+ importance_sample_ratio=cfg.MODEL.MASK_FORMER.IMPORTANCE_SAMPLE_RATIO,
+ )
+
+ return {
+ "cfg": cfg,
+ "backbone": backbone,
+ "sem_seg_head": sem_seg_head,
+ "criterion_2d": criterion_2d,
+ "criterion_3d": criterion_3d,
+ "num_queries": cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES,
+ "object_mask_threshold": cfg.MODEL.MASK_FORMER.TEST.OBJECT_MASK_THRESHOLD,
+ "overlap_threshold": cfg.MODEL.MASK_FORMER.TEST.OVERLAP_THRESHOLD,
+ "metadata": MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
+ "size_divisibility": cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY,
+ "sem_seg_postprocess_before_inference": (
+ cfg.MODEL.MASK_FORMER.TEST.SEM_SEG_POSTPROCESSING_BEFORE_INFERENCE
+ or cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON
+ or cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON
+ ),
+ "pixel_mean": cfg.MODEL.PIXEL_MEAN,
+ "pixel_std": cfg.MODEL.PIXEL_STD,
+ # inference
+ "semantic_on": cfg.MODEL.MASK_FORMER.TEST.SEMANTIC_ON,
+ "instance_on": cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON,
+ "panoptic_on": cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON,
+ "test_topk_per_image": cfg.TEST.DETECTIONS_PER_IMAGE,
+ }
+
+ @property
+ def device(self):
+ return self.pixel_mean.device
+
+ def forward(self, batched_inputs):
+ """
+ Args:
+ batched_inputs: a list, batched outputs of :class:`DatasetMapper`.
+ Each item in the list contains the inputs for one image.
+ For now, each item in the list is a dict that contains:
+ * "image": Tensor, image in (C, H, W) format.
+ * "instances": per-region ground truth
+ * Other information that's included in the original dicts, such as:
+ "height", "width" (int): the output resolution of the model (may be different
+ from input resolution), used in inference.
+ Returns:
+ list[dict]:
+ each dict has the results for one image. The dict contains the following keys:
+
+ * "sem_seg":
+ A Tensor that represents the
+ per-pixel segmentation prediced by the head.
+ The prediction has shape KxHxW that represents the logits of
+ each class for each pixel.
+ * "panoptic_seg":
+ A tuple that represent panoptic output
+ panoptic_seg (Tensor): of shape (height, width) where the values are ids for each segment.
+ segments_info (list[dict]): Describe each segment in `panoptic_seg`.
+ Each dict contains keys "id", "category_id", "isthing".
+ """
+ ## make new images
+ batched_inputs_new = []
+ for batched_input in batched_inputs:
+ ori_infos = {"height": batched_input["height"],
+ "width": batched_input["width"], 
+ "image": batched_input["image"],
+ # "file_name": batched_input["file_name"],
+ }
+ if "instances" in batched_input.keys():
+ ori_instances = batched_input["instances"]
+ ori_instances.original_indices = torch.arange(0, len(ori_instances)).long()
+ ori_infos["instances"] = ori_instances
+ batched_inputs_new.append(ori_infos)
+ ## cropped patches
+ # pdb.set_trace()
+ crop_region = batched_input["crop_region"]
+ crop_images = batched_input["image_crop"]
+ crop_o_width = int(crop_region[0][2]-crop_region[0][0])
+ crop_o_height = int(crop_region[0][3]-crop_region[0][1])
+
+ if "instances_crop" in batched_input.keys():
+ crop_instances = batched_input["instances_crop"]
+ else:
+ crop_instances = None
+
+ for crop_index, crop_image in enumerate(crop_images):
+ crop_infos = {"height": crop_o_height, "width": crop_o_width, "image": crop_image}
+ if not crop_instances == None:
+ crop_instance = crop_instances[crop_index]
+ crop_instance.original_indices = torch.arange(0, len(crop_instance)).long()
+ crop_infos["instances"] = crop_instance
+ batched_inputs_new.append(crop_infos)
+
+ images = [x["image"].to(self.device) for x in batched_inputs_new]
+ ## +1 means 
+ num_views = self.cfg.ENTITY.CROP_SAMPLE_NUM_TRAIN+1 if self.training else self.cfg.ENTITY.CROP_SAMPLE_NUM_TEST+1
+ for i in range(len(images)):
+ if i%num_views==0:
+ continue
+ _, c_h, c_w = images[i].shape
+ if "instances" in batched_inputs_new[i].keys():
+ batched_inputs_new[i]["instances"]._image_size = (c_h, c_w)
+ 
+ images = [(x - self.pixel_mean) / self.pixel_std for x in images]
+ images = ImageList.from_tensors(images, self.size_divisibility)
+
+ features = self.backbone(images.tensor)
+ outputs_2d, outputs_3d = self.sem_seg_head(features)
+
+ if self.training:
+ if self.cfg.ENTITY.ENABLE:
+ for i in range(len(batched_inputs_new)):
+ batched_inputs_new[i]["instances"].gt_classes[:] = 0
+ 
+ if "instances" in batched_inputs[0]:
+ gt_instances = [x["instances"].to(self.device) for x in batched_inputs_new]
+ targets_2d = self.prepare_targets_2d(copy.deepcopy(gt_instances), copy.deepcopy(images))
+ targets_3d = self.prepare_targets_3d(copy.deepcopy(gt_instances), copy.deepcopy(images), num_views)
+ else:
+ targets = None
+
+ # bipartite matching-based loss
+ losses = {}
+ losses_2d = self.criterion_2d(outputs_2d, targets_2d)
+ losses_3d = self.criterion_3d(outputs_3d, targets_3d)
+
+ for k in list(losses_2d.keys()):
+ if k in self.criterion_2d.weight_dict:
+ losses[k+"_2d"] = losses_2d[k] * self.criterion_2d.weight_dict[k] * 0.5
+ else:
+ # remove this loss if not specified in `weight_dict`
+ losses_2d.pop(k)
+ 
+ for k in list(losses_3d.keys()):
+ if k in self.criterion_3d.weight_dict:
+ losses[k+"_3d"] = losses_3d[k] * self.criterion_3d.weight_dict[k]
+ else:
+ # remove this loss if not specified in `weight_dict`
+ losses_3d.pop(k)
+ return losses
+ else:
+ mask_cls_results_3d = outputs_3d["pred_logits"][0] ## 100,2
+ mask_pred_results_3d = outputs_3d["pred_masks"][0] ## 100,5,200, 304
+
+ mask_cls_results_2d = outputs_2d["pred_logits"]
+ mask_pred_results_2d = outputs_2d["pred_masks"]
+ # upsample masks
+ 
+ mask_pred_results_3d = retry_if_cuda_oom(F.interpolate)(
+ mask_pred_results_3d,
+ size=(images.tensor.shape[-2], images.tensor.shape[-1]),
+ mode="bilinear",
+ align_corners=False,
+ )
+
+ mask_pred_results_2d = F.interpolate(
+ mask_pred_results_2d,
+ size=(images.tensor.shape[-2], images.tensor.shape[-1]),
+ mode="bilinear",
+ align_corners=False,
+ )
+
+ del outputs_2d, outputs_3d
+ 
+ crop_regions = batched_input["crop_region"][:num_views-1]
+ processed_results = retry_if_cuda_oom(self.inference_whole_views)(
+ mask_cls_results_3d,
+ mask_pred_results_3d,
+ mask_cls_results_2d,
+ mask_pred_results_2d,
+ batched_inputs_new,
+ images.image_sizes, 
+ crop_regions)
+
+ # processed_results = retry_if_cuda_oom(self.instance_inference_nonoverlap)(
+ # mask_cls_results_2d[0], 
+ # mask_pred_results_2d[0],
+ # batched_inputs_new[0], 
+ # images.image_sizes[0])
+
+ return [{"instances": processed_results}]
+
+ def prepare_targets_2d(self, targets, images):
+ h_pad, w_pad = images.tensor.shape[-2:]
+ new_targets = []
+ for targets_per_image in targets:
+ gt_masks = targets_per_image.gt_masks.tensor 
+ gt_valid = targets_per_image.gt_boxes_valid
+ padded_masks = torch.zeros((gt_masks.shape[0], h_pad, w_pad), dtype=gt_masks.dtype, device=gt_masks.device)
+ padded_masks[:, : gt_masks.shape[1], : gt_masks.shape[2]] = gt_masks
+ valid_index = torch.nonzero(gt_valid).flatten()
+ new_targets.append(
+ {
+ "labels": targets_per_image.gt_classes[valid_index],
+ "masks": padded_masks[valid_index],
+ }
+ )
+ return new_targets
+ 
+ def prepare_targets_3d(self, targets_ori, images, num_views):
+ T = num_views
+ B = int(len(targets_ori) / T)
+ h_pad, w_pad = images.tensor.shape[-2:]
+ 
+ ## reshape to new targets
+ new_targets = []
+ for count, target in enumerate(targets_ori):
+ b_index, t_index = int(count // T), int(count % T)
+ if t_index == 0:
+ new_targets.append([target])
+ else:
+ new_targets[b_index].append(target)
+
+ gt_instances = []
+ for count, targets in enumerate(new_targets):
+ _num_instance = len(targets[0])
+ mask_shape = [_num_instance, T, h_pad, w_pad]
+ gt_masks_per_view = torch.zeros(mask_shape, dtype=torch.bool, device=self.device)
+
+ for v_i, targets_per_view in enumerate(targets):
+ assert torch.all(targets[0].original_indices == targets_per_view.original_indices)
+ 
+ gt_ids_per_view = []
+ gt_ids_per_valid = []
+ gt_ids_categories = []
+ ## view first, then entities
+ for v_i, targets_per_view in enumerate(targets):
+ targets_per_view = targets_per_view.to(self.device)
+ h, w = targets_per_view.image_size
+ for i_i, (instance_mask, instance_valid) in enumerate(zip(targets_per_view.gt_masks.tensor, targets_per_view.gt_boxes_valid)):
+ if instance_valid == 1:
+ gt_masks_per_view[i_i, v_i, :h, :w] = instance_mask
+ gt_ids_per_valid.append(targets_per_view.gt_boxes_valid[None,:])
+ gt_ids_per_view.append(targets_per_view.original_indices[None,:])
+ gt_ids_categories.append(targets_per_view.gt_classes[None, :])
+ ## (num_instances, num_views)
+ gt_ids_per_valid = torch.cat(gt_ids_per_valid, dim=0).permute((1,0))
+ gt_ids_per_view = torch.cat(gt_ids_per_view, dim=0).permute((1,0))
+ gt_ids_categories = torch.cat(gt_ids_categories, dim=0).permute((1,0))
+ 
+ gt_ids_per_view[gt_ids_per_valid == 0] = -1
+ valid_idx = (gt_ids_per_view != 1).any(dim=-1)
+ ## categoreis 
+ gt_classes_per_group = gt_ids_categories[:,0] ## N
+ gt_ids_per_group = gt_ids_per_view ## N, num_views
+ gt_masks_per_group = gt_masks_per_view.float() ## N, num_views, H, W
+
+ ## 
+ gt_instances.append({"labels": gt_classes_per_group, 
+ "ids": gt_ids_per_group,
+ "masks": gt_masks_per_group})
+
+ return gt_instances
+
+ def semantic_inference(self, mask_cls, mask_pred):
+ mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1]
+ mask_pred = mask_pred.sigmoid()
+ semseg = torch.einsum("qc,qhw->chw", mask_cls, mask_pred)
+ return semseg
+
+ def panoptic_inference(self, mask_cls, mask_pred):
+ scores, labels = F.softmax(mask_cls, dim=-1).max(-1)
+ mask_pred = mask_pred.sigmoid()
+
+ keep = labels.ne(self.sem_seg_head.num_classes) & (scores > self.object_mask_threshold)
+ cur_scores = scores[keep]
+ cur_classes = labels[keep]
+ cur_masks = mask_pred[keep]
+ cur_mask_cls = mask_cls[keep]
+ cur_mask_cls = cur_mask_cls[:, :-1]
+
+ cur_prob_masks = cur_scores.view(-1, 1, 1) * cur_masks
+
+ h, w = cur_masks.shape[-2:]
+ panoptic_seg = torch.zeros((h, w), dtype=torch.int32, device=cur_masks.device)
+ segments_info = []
+
+ current_segment_id = 0
+
+ if cur_masks.shape[0] == 0:
+ # We didn't detect any mask :(
+ return panoptic_seg, segments_info
+ else:
+ # take argmax
+ cur_mask_ids = cur_prob_masks.argmax(0)
+ stuff_memory_list = {}
+ for k in range(cur_classes.shape[0]):
+ pred_class = cur_classes[k].item()
+ isthing = pred_class in self.metadata.thing_dataset_id_to_contiguous_id.values()
+ mask_area = (cur_mask_ids == k).sum().item()
+ original_area = (cur_masks[k] >= 0.5).sum().item()
+ mask = (cur_mask_ids == k) & (cur_masks[k] >= 0.5)
+
+ if mask_area > 0 and original_area > 0 and mask.sum().item() > 0:
+ if mask_area / original_area < self.overlap_threshold:
+ continue
+
+ # merge stuff regions
+ if not isthing:
+ if int(pred_class) in stuff_memory_list.keys():
+ panoptic_seg[mask] = stuff_memory_list[int(pred_class)]
+ continue
+ else:
+ stuff_memory_list[int(pred_class)] = current_segment_id + 1
+
+ current_segment_id += 1
+ panoptic_seg[mask] = current_segment_id
+
+ segments_info.append(
+ {
+ "id": current_segment_id,
+ "isthing": bool(isthing),
+ "category_id": int(pred_class),
+ }
+ )
+ return panoptic_seg, segments_info
+ 
+ def instance_inference_nonoverlap(self, mask_cls, mask_pred):
+ # mask_pred is already processed to have the same shape as original input
+ image_size = mask_pred.shape[-2:]
+
+ # [Q, K]
+ scores = F.softmax(mask_cls, dim=-1)[:, :-1]
+ labels = torch.arange(self.sem_seg_head.num_classes, device=self.device).unsqueeze(0).repeat(self.num_queries, 1).flatten(0, 1)
+ # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
+ scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.test_topk_per_image, sorted=False)
+ labels_per_image = labels[topk_indices]
+
+ topk_indices = topk_indices // self.sem_seg_head.num_classes
+ # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
+ mask_pred = mask_pred[topk_indices]
+
+ ###### ranks
+ pred_masks = (mask_pred>0).float()
+ pred_masks_logits = mask_pred.sigmoid()
+ pred_scores = scores_per_image
+
+ _, m_H, m_W = pred_masks.shape
+ mask_id = torch.zeros((m_H, m_W), dtype=torch.int).to(pred_masks.device)
+ sorted_scores, ranks = torch.sort(pred_scores)
+ ranks = ranks + 1
+ for index in ranks:
+ mask_id[(pred_masks[index-1]==1)] = int(index)
+ # re-generate mask
+ new_scores = []
+ new_masks = []
+ new_masks_logits = []
+ entity_nums = len(ranks)
+ for ii in range(entity_nums):
+ index = int(ranks[entity_nums-ii-1])
+ score = sorted_scores[entity_nums-ii-1]
+ new_scores.append(score)
+ new_masks.append((mask_id==index).float())
+ new_masks_logits.append(pred_masks_logits[index-1])
+ 
+ new_scores = torch.stack(new_scores)
+ new_masks = torch.stack(new_masks)
+ new_masks_logits = torch.stack(new_masks_logits)
+
+ result = Instances(image_size)
+ # mask (before sigmoid)
+ result.pred_masks = new_masks
+ result.pred_boxes = Boxes(torch.zeros(new_masks.size(0), 4))
+ # Uncomment the following to get boxes from masks (this is slow)
+
+ # calculate average mask prob
+ mask_scores_per_image = (new_masks_logits.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
+ result.scores = new_scores * mask_scores_per_image
+ result.pred_classes = labels_per_image
+ return result
+
+ def instance_inference(self, mask_cls, mask_pred):
+ # mask_pred is already processed to have the same shape as original input
+ image_size = mask_pred.shape[-2:]
+ 
+ # [Q, K]
+ scores = F.softmax(mask_cls, dim=-1)[:, :-1]
+ labels = torch.arange(self.sem_seg_head.num_classes, device=self.device).unsqueeze(0).repeat(self.num_queries, 1).flatten(0, 1)
+ # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
+ scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.test_topk_per_image, sorted=False)
+ labels_per_image = labels[topk_indices]
+ 
+ topk_indices = topk_indices // self.sem_seg_head.num_classes
+ # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
+ mask_pred = mask_pred[topk_indices]
+
+ # if this is panoptic segmentation, we only keep the "thing" classes
+ if self.panoptic_on:
+ keep = torch.zeros_like(scores_per_image).bool()
+ for i, lab in enumerate(labels_per_image):
+ keep[i] = lab in self.metadata.thing_dataset_id_to_contiguous_id.values()
+
+ scores_per_image = scores_per_image[keep]
+ labels_per_image = labels_per_image[keep]
+ mask_pred = mask_pred[keep]
+
+ result = Instances(image_size)
+ # mask (before sigmoid)
+ result.pred_masks = (mask_pred > 0).float()
+ result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
+ # Uncomment the following to get boxes from masks (this is slow)
+ # result.pred_boxes = BitMasks(mask_pred > 0).get_bounding_boxes()
+
+ # calculate average mask prob
+ mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
+ # pdb.set_trace()
+ result.scores = scores_per_image * mask_scores_per_image
+ result.pred_classes = labels_per_image
+ return result
+ 
+ def inference_whole_views(self, pred_cls, pred_masks, pred_cls_2d, pred_masks_2d, batched_inputs, image_sizes, crop_regions):
+ ## pred_masks: [100, 5, 800, 1216]
+ ## pred_masks_2d: [5, 100, 800, 1216]
+ scores = F.softmax(pred_cls, dim=-1)[:,:-1] # 100,1
+ scores_2d = F.softmax(pred_cls_2d, dim=-1)[:, :, :-1] # 5, 100, 1
+ 
+ # scores = (scores+scores_2d[0])/2
+ labels = torch.arange(self.sem_seg_head.num_classes, device=self.device).unsqueeze(0).repeat(self.num_queries, 1).flatten(0, 1)
+ ### keep all the indices
+ scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
+ labels_per_image = labels[topk_indices]
+ # topk_indices = topk_indices // self.sem_seg_head.num_classes
+ topk_indices = torch.div(topk_indices, self.sem_seg_head.num_classes, rounding_mode="trunc")
+ pred_masks = pred_masks[topk_indices]
+ pred_masks = pred_masks.permute((1,0,2,3))
+
+ new_pred_masks = []
+ for view_index, (pred_masks_per_view, batched_input_per_view, image_size_per_view) in enumerate(zip(pred_masks, batched_inputs, image_sizes)):
+ O_H = batched_input_per_view["height"]
+ O_W = batched_input_per_view["width"]
+
+ SO_H, SO_W = image_size_per_view
+
+ pred_masks_per_view = pred_masks_per_view[..., : SO_H, :SO_W]
+ pred_masks_per_view = F.interpolate(pred_masks_per_view[None], size=(O_H, O_W), mode="bilinear", align_corners=False)
+
+ new_pred_masks.append(pred_masks_per_view[0].sigmoid())
+ 
+ ## fuse the masks
+ full_image_masks = new_pred_masks[0]
+ 
+ ## fuse crop image
+ fused_image_masks = torch.zeros_like(full_image_masks).float()
+ fused_image_masks_valid = torch.zeros_like(full_image_masks).float() + 1e-16
+ for crop_region_per_view, pred_masks_per_view in zip(crop_regions, new_pred_masks[1:]):
+ x0, y0, x1, y1 = crop_region_per_view
+ fused_image_masks[..., y0:y1, x0:x1] += pred_masks_per_view
+ fused_image_masks_valid[..., y0:y1, x0:x1] += 1
+ 
+ # add original masks
+ fused_image_masks += full_image_masks
+ fused_image_masks_valid += 1
+
+ ## average
+ fuse_image_masks = fused_image_masks / fused_image_masks_valid
+
+ ###### change to the single image, begin to non_overlap_supression
+ ## ranks
+ pred_masks_logits = fuse_image_masks
+ pred_masks = (fuse_image_masks>0.5).float()
+ pred_scores = scores_per_image
+
+ _, m_H, m_W = pred_masks.shape
+ ## for visualization
+ mask_id = torch.zeros((m_H, m_W), dtype=torch.int).to(pred_masks.device)
+ 
+ # mask_id_colors = np.zeros((m_H, m_W, 3), dtype=np.uint8)
+ # pred_masks_np = pred_masks.cpu().numpy()
+
+ sorted_scores, ranks = torch.sort(pred_scores)
+ ranks = ranks + 1
+ for index in ranks:
+ mask_id[(pred_masks[index-1]==1)] = int(index)
+ # mask_id_colors[(pred_masks_np[index-1]==1)] = self.colors[index]
+ # base_path = "/group/20018/gavinqi/vis_entityv2_release_debug"
+ # pdb.set_trace()
+ # file_name = batched_inputs[0]["file_name"]
+ # split_index, img_name = file_name.split("/")[-2:]
+ # save_name = img_name.split(".")[0]+".png"
+ # if not os.path.exists(os.path.join(base_path, save_name)):
+ # cv2.imwrite(os.path.join(base_path, save_name), mask_id_colors)
+ # re-generate mask
+ new_scores = []
+ new_masks = []
+ new_masks_logits = []
+ entity_nums = len(ranks)
+ for ii in range(entity_nums):
+ index = int(ranks[entity_nums-ii-1])
+ score = sorted_scores[entity_nums-ii-1]
+ new_scores.append(score)
+ new_masks.append((mask_id==index).float())
+ new_masks_logits.append(pred_masks_logits[index-1])
+ 
+ new_scores = torch.stack(new_scores)
+ new_masks = torch.stack(new_masks)
+ new_masks_logits = torch.stack(new_masks_logits)
+ # make result
+ image_size = (batched_inputs[0]["height"], batched_inputs[0]["width"])
+ result = Instances(image_size)
+ # mask (before sigmoid)
+ result.pred_masks = new_masks
+ result.pred_boxes = Boxes(torch.zeros(new_masks.size(0), 4))
+ # Uncomment the following to get boxes from masks (this is slow)
+
+ # calculate average mask prob
+ mask_scores_per_image = (new_masks_logits.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
+ result.scores = new_scores * mask_scores_per_image
+ result.pred_classes = labels_per_image
+ return result

annotator/entityseg/mask2former/data/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Facebook, Inc. and its affiliates.

annotator/entityseg/mask2former/data/dataset_mappers/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Facebook, Inc. and its affiliates.

annotator/entityseg/mask2former/data/dataset_mappers/crop_augmentations.py

@@ -0,0 +1,421 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+"""
+Implement many useful :class:`Augmentation`.
+"""
+import numpy as np
+import sys
+from typing import Tuple
+from PIL import Image
+import random
+
+from fvcore.transforms.transform import NoOpTransform, Transform
+
+from detectron2.data.transforms.augmentation import Augmentation
+import pdb
+import math
+
+import logging
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+from PIL import Image
+from collections import defaultdict
+import copy
+from detectron2.data import transforms as T
+from detectron2.structures import (
+ BitMasks,
+ Boxes,
+ BoxMode,
+ Instances,
+ Keypoints,
+ PolygonMasks,
+ RotatedBoxes,
+ polygons_to_bitmask,
+)
+from detectron2.utils.file_io import PathManager
+
+__all__ = [
+ "BatchResizeShortestEdge",
+ "EntityCrop",
+]
+
+class BatchResizeTransform(Transform):
+ """
+ Resize the image to a target size.
+ """
+
+ def __init__(self, h, w, new_h, new_w, 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, imgs, interp=None):
+ dim_num = len(imgs.shape)
+ assert dim_num == 4
+ interp_method = interp if interp is not None else self.interp
+ resized_imgs = []
+ for img in imgs:
+ if len(img.shape) > 2 and img.shape[2] == 1:
+ pil_image = Image.fromarray(img[:, :, 0], mode="L")
+ else:
+ pil_image = Image.fromarray(img)
+ pil_image = pil_image.resize((self.new_w, self.new_h), interp_method)
+ ret = np.asarray(pil_image)
+ if len(img.shape) > 2 and img.shape[2] == 1:
+ ret = np.expand_dims(ret, -1)
+ resized_imgs.append(ret)
+ resized_imgs = np.stack(resized_imgs)
+ return resized_imgs
+
+ def apply_coords(self, coords):
+ coords[:, 0] = coords[:, 0] * (self.new_w * 1.0 / self.w)
+ coords[:, 1] = coords[:, 1] * (self.new_h * 1.0 / self.h)
+ return coords
+ 
+ def apply_box(self, boxes):
+ boxes = boxes[0]
+ new_boxes = super(BatchResizeTransform, self).apply_box(boxes[:,:4])
+ boxes[...,:4] = new_boxes
+ return boxes[None]
+
+ def apply_segmentation(self, segmentation):
+ if len(segmentation.shape)==3:
+ segmentation = segmentation[..., None]
+ segmentation = self.apply_image(segmentation, interp=Image.NEAREST)
+ segmentation = segmentation[..., 0]
+ else:
+ segmentation = self.apply_image(segmentation, interp=Image.NEAREST)
+ return segmentation
+
+class EntityCropTransform(Transform):
+ """
+ Consectively crop the images
+ """
+ def __init__(self, crop_axises, crop_indexes):
+ super().__init__()
+ self._set_attributes(locals())
+ 
+ def apply_image(self, img):
+ """
+ Args:
+ img (ndarray): of shape NxHxWxC, or HxWxC or HxW. The array can be
+ of type uint8 in range [0, 255], or floating point in range
+ [0, 1] or [0, 255]
+ returns:
+ ndarray: cropped images
+ """
+ dim_num = len(img.shape)
+ imgs = []
+ 
+ for crop_axis in self.crop_axises:
+ x0, y0, x1, y1 = crop_axis
+ if dim_num <= 3:
+ crop_img = img[y0:y1, x0:x1]
+ else:
+ crop_img = img[..., y0:y1, x0:x1, :]
+ imgs.append(crop_img)
+
+ if dim_num <= 3:
+ imgs = np.stack(imgs, axis=0)
+ else:
+ imgs = np.concatenate(imgs, axis=0)
+ return imgs
+ 
+ def apply_coords(self, coords: np.ndarray, x0, y0):
+ coords[:, 0] -= x0
+ coords[:, 1] -= y0
+ return coords
+ 
+ def apply_box(self, box: np.ndarray) -> np.ndarray:
+ """
+ box: Nx4, [x0, y0, x1, y1]
+ """
+ idxs = np.array([(0, 1), (2, 1), (0, 3), (2, 3)]).flatten()
+ coords = np.asarray(box).reshape(-1, 4)[:, idxs].reshape(-1, 2)
+ split_boxes = []
+ crop_ws, crop_hs = [], []
+ for crop_axis in self.crop_axises:
+ startw, starth, endw, endh = crop_axis
+ coords_new = self.apply_coords(copy.deepcopy(coords), startw, starth).reshape((-1, 4, 2))
+ minxy = coords_new.min(axis=1)
+ maxxy = coords_new.max(axis=1)
+ trans_boxes = np.concatenate((minxy, maxxy), axis=1)
+ 
+ crop_ws.append(endw-startw)
+ crop_hs.append(endh-starth)
+ split_boxes.append(trans_boxes)
+ split_boxes = np.stack(split_boxes, axis=1)
+ ### clip to the image boundary
+ ## assert each crop size is equal
+ for crop_index, (crop_w, crop_h) in enumerate(zip(crop_ws, crop_hs)):
+ assert crop_w == crop_ws[0], "crop width is not equal, crop_{}: {}, crop_0: {}".format(crop_index, crop_w, crop_ws[0])
+ assert crop_h == crop_hs[0], "crop height is not equal, crop_{}: {}, crop_0: {}".format(crop_index, crop_h, crop_hs[0])
+ crop_w = crop_ws[0]
+ crop_h = crop_hs[0]
+ # pdb.set_trace()
+ split_boxes[...,0::2] = np.clip(split_boxes[...,0::2], 0, crop_w)
+ split_boxes[...,1::2] = np.clip(split_boxes[...,1::2], 0, crop_h)
+ valid_inds = (split_boxes[...,2]>split_boxes[...,0]) & (split_boxes[...,3]>split_boxes[...,1])
+ split_infos = np.concatenate((split_boxes, valid_inds[...,None]), axis=-1)
+ return split_infos
+
+class BatchResizeShortestEdge(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, short_edge_length, max_size=sys.maxsize, sample_style="range", interp=Image.BILINEAR
+ ):
+ """
+ Args:
+ short_edge_length (list[int]): If ``sample_style=="range"``,
+ a [min, max] interval from which to sample the shortest edge length.
+ If ``sample_style=="choice"``, a list of shortest edge lengths to sample from.
+ max_size (int): maximum allowed longest edge length.
+ sample_style (str): either "range" or "choice".
+ """
+ super().__init__()
+ assert sample_style in ["range", "choice"], sample_style
+
+ self.is_range = sample_style == "range"
+ if isinstance(short_edge_length, int):
+ short_edge_length = (short_edge_length, short_edge_length)
+ if self.is_range:
+ assert len(short_edge_length) == 2, (
+ "short_edge_length must be two values using 'range' sample style."
+ f" Got {short_edge_length}!"
+ )
+ self._init(locals())
+
+ def get_transform(self, image):
+ dim_num = len(image.shape)
+ assert dim_num == 4, "the tensor should be in [B, H, W, C]"
+ h, w = image.shape[1:3]
+ if self.is_range:
+ size = np.random.randint(self.short_edge_length[0], self.short_edge_length[1] + 1)
+ else:
+ size = np.random.choice(self.short_edge_length)
+ if size == 0:
+ return NoOpTransform()
+
+ scale = size * 1.0 / min(h, w)
+ if h < w:
+ newh, neww = size, scale * w
+ else:
+ newh, neww = scale * h, size
+ if max(newh, neww) > self.max_size:
+ scale = self.max_size * 1.0 / max(newh, neww)
+ newh = newh * scale
+ neww = neww * scale
+ neww = int(neww + 0.5)
+ newh = int(newh + 0.5)
+ return BatchResizeTransform(h, w, newh, neww, self.interp)
+
+class EntityCrop(Augmentation):
+ def __init__(self, crop_ratio, stride_ratio, sample_num, is_train):
+ super().__init__()
+ self._init(locals())
+
+ def get_transform(self, image):
+ h, w = image.shape[:2]
+ crop_axises, crop_indexes = self.get_crop_axises((h, w))
+ transform = EntityCropTransform(crop_axises, crop_indexes)
+ return transform
+ 
+ def get_crop_axises(self, image_size):
+ h, w = image_size
+ crop_w = int(self.crop_ratio*w)
+ crop_h = int(self.crop_ratio*h)
+ # if self.is_train:
+ stride_w = int(self.stride_ratio*w)
+ stride_h = int(self.stride_ratio*h)
+ # pdb.set_trace()
+
+ crop_axises = []
+ for starth in range(0, h, stride_h):
+ for startw in range(0, w, stride_w):
+ endh = min(starth+crop_h, h)
+ endw = min(startw+crop_w, w)
+ starth = int(endh-crop_h)
+ startw = int(endw-crop_w)
+ crop_axises.append([startw, starth, endw, endh])
+ if self.is_train:
+ crop_indexes = random.sample([i for i in range(len(crop_axises))], self.sample_num)
+ crop_axises = [crop_axises[i] for i in crop_indexes]
+ else:
+ crop_indexes = [i for i in range(self.sample_num)]
+ # left_upper = [0, 0, crop_w, crop_h]
+ # right_upper = [w-crop_w, 0, w, crop_h]
+ # left_bottom = [0, h-crop_h, crop_w, h]
+ # right_bottom = [w-crop_w, h-crop_h, w, h]
+ 
+ # crop_axises = [left_upper, right_upper, left_bottom, right_bottom]
+ # crop_indexes = [0,1,2,3]
+ assert len(crop_axises)==len(crop_indexes)
+ return crop_axises, crop_indexes
+
+def transform_instance_annotations_crop(
+ annotation, transforms, image_size, *, keypoint_hflip_indices=None
+):
+ """
+ Apply transforms to box, segmentation and keypoints annotations of a single instance.
+
+ It will use `transforms.apply_box` for the box, and
+ `transforms.apply_coords` for segmentation polygons & keypoints.
+ If you need anything more specially designed for each data structure,
+ you'll need to implement your own version of this function or the transforms.
+
+ Args:
+ annotation (dict): dict of instance annotations for a single instance.
+ It will be modified in-place.
+ transforms (TransformList or list[Transform]):
+ image_size (tuple): the height, width of the transformed image
+ keypoint_hflip_indices (ndarray[int]): see `create_keypoint_hflip_indices`.
+
+ Returns:
+ dict:
+ the same input dict with fields "bbox", "segmentation", "keypoints"
+ transformed according to `transforms`.
+ The "bbox_mode" field will be set to XYXY_ABS.
+ """
+ if isinstance(transforms, (tuple, list)):
+ transforms = T.TransformList(transforms)
+ # bbox is 1d (per-instance bounding box)
+ bbox = BoxMode.convert(annotation["bbox"], annotation["bbox_mode"], BoxMode.XYXY_ABS)
+ 
+ # clip transformed bbox to image size
+ bboxes_info = transforms.apply_box(np.array([bbox]))[0].clip(min=0)
+ annotation["bbox"] = np.minimum(bbox, list(image_size + image_size)[::-1])
+ annotation["bbox"] = bboxes_info[...,:4]
+ annotation["bbox_mode"] = BoxMode.XYXY_ABS
+ annotation["bbox_valid"] = bboxes_info[...,4]
+ for transform_type in transforms:
+ if isinstance(transform_type, EntityCropTransform):
+ annotation["crop_axises"] = transform_type.crop_axises
+ annotation["crop_indexes"] = transform_type.crop_indexes
+
+ if "segmentation" in annotation:
+ segm = annotation["segmentation"]
+ assert isinstance(segm, dict), "requiring segmentation encoding -> RLE"
+ # RLE
+ mask = mask_util.decode(segm)
+ mask = transforms.apply_segmentation(mask)
+ annotation["segmentation"] = mask
+ return annotation
+
+def annotations_to_instances_crop(annos, image_size, mask_format="polygon", return_indexes=False):
+ """
+ Create an :class:`Instances` object used by the models,
+ from instance annotations in the dataset dict.
+
+ Args:
+ annos (list[dict]): a list of instance annotations in one image, each
+ element for one instance.
+ image_size (tuple): height, width
+
+ Returns:
+ Instances:
+ It will contain fields "gt_boxes", "gt_classes",
+ "gt_masks", "gt_keypoints", if they can be obtained from `annos`.
+ This is the format that builtin models expect.
+ """
+ ###
+ all_boxes = []
+ all_boxes_valid = []
+ all_classes = []
+ all_segmentations = []
+ all_iscrowds = []
+ # pdb.set_trace()
+ annos_num = len(annos)
+ patches_num = len(annos[0]["bbox"])
+ for ann_index, obj in enumerate(annos):
+ for split_index in range(len(obj["bbox"])):
+ all_boxes.append(BoxMode.convert(obj["bbox"][split_index], obj["bbox_mode"], BoxMode.XYXY_ABS))
+ all_boxes_valid.append(obj["bbox_valid"][split_index])
+ all_classes.append(obj["category_id"])
+ all_segmentations.append(obj["segmentation"][split_index])
+ all_iscrowds.append(obj["iscrowd"])
+ # print("ann_index:{}, split_index:{}".format(ann_index, split_index))
+ 
+ new_targets = []
+ crop_axises = annos[0]["crop_axises"]
+ # pdb.set_trace()
+ crop_size = (crop_axises[0][3], crop_axises[0][2])
+ crop_axises = torch.tensor(crop_axises)
+ 
+ for split_index in range(patches_num):
+ new_targets.append(Instances(crop_size))
+ # pdb.set_trace()
+ ## boxes
+ new_targets[-1].gt_boxes = Boxes(all_boxes[split_index::patches_num])
+ new_targets[-1].gt_boxes_valid = torch.tensor(all_boxes_valid[split_index::patches_num], dtype=torch.int64)
+ ## categories
+ new_targets[-1].gt_classes = torch.tensor(all_classes[split_index::patches_num], dtype=torch.int64)
+
+ ## masks
+ if "segmentation" in annos[0]:
+ new_targets[-1].gt_masks = BitMasks(torch.stack([torch.from_numpy(np.ascontiguousarray(x)) for x in all_segmentations[split_index::patches_num]]))
+ 
+ # pdb.set_trace()
+ if return_indexes:
+ return new_targets, crop_axises, annos[0]["crop_indexes"]
+ else:
+ return new_targets, crop_axises
+
+class EntityCascadedCrop(Augmentation):
+ def __init__(self, crop_ratio, stride_ratio, sample_num, cascade_num, is_train):
+ super().__init__()
+ self._init(locals())
+
+ def get_transform(self, image):
+ h, w = image.shape[:2]
+ crop_axises, crop_indexes = self.get_crop_axises((h, w))
+ transform = EntityCropTransform(crop_axises, crop_indexes)
+ return transform
+ 
+ def get_crop_axises(self, image_size):
+ h, w = image_size
+ # for i in range(self.cascade_num):
+ # crop_w = int((self.crop_ratio**(i+1))*w)
+ # crop_h = int((self.crop_ratio**(i+1))*h)
+ # stride_w = int((self.stride_ratio**(i+1))*w)
+ # stride_h = int((self.stride_ratio**(i+1))*h)
+ # crop_axises = []
+ # if i==0:
+
+ # for starth in range(0, )
+
+
+ crop_axises = []
+ for starth in range(0, h, stride_h):
+ for startw in range(0, w, stride_w):
+ endh = min(starth+crop_h, h)
+ endw = min(startw+crop_w, w)
+ starth = int(endh-crop_h)
+ startw = int(endw-crop_w)
+ crop_axises.append([startw, starth, endw, endh])
+ if self.is_train:
+ crop_indexes = random.sample([i for i in range(len(crop_axises))], self.sample_num)
+ crop_axises = [crop_axises[i] for i in crop_indexes]
+ else:
+ crop_indexes = [i for i in range(self.sample_num)]
+ # left_upper = [0, 0, crop_w, crop_h]
+ # right_upper = [w-crop_w, 0, w, crop_h]
+ # left_bottom = [0, h-crop_h, crop_w, h]
+ # right_bottom = [w-crop_w, h-crop_h, w, h]
+ 
+ # crop_axises = [left_upper, right_upper, left_bottom, right_bottom]
+ # crop_indexes = [0,1,2,3]
+ assert len(crop_axises)==len(crop_indexes)
+ return crop_axises, crop_indexes

annotator/entityseg/mask2former/maskformer_model.py

@@ -0,0 +1,446 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from typing import Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from detectron2.config import configurable
+from detectron2.data import MetadataCatalog
+from detectron2.modeling import META_ARCH_REGISTRY, build_backbone, build_sem_seg_head
+from detectron2.modeling.backbone import Backbone
+from detectron2.modeling.postprocessing import sem_seg_postprocess
+from detectron2.structures import Boxes, ImageList, Instances, BitMasks
+from detectron2.utils.memory import retry_if_cuda_oom
+
+from .modeling.criterion import SetCriterion
+from .modeling.matcher import HungarianMatcher
+
+
+@META_ARCH_REGISTRY.register()
+class MaskFormer(nn.Module):
+ """
+ Main class for mask classification semantic segmentation architectures.
+ """
+
+ @configurable
+ def __init__(
+ self,
+ *,
+ cfg,
+ backbone: Backbone,
+ sem_seg_head: nn.Module,
+ criterion: nn.Module,
+ num_queries: int,
+ object_mask_threshold: float,
+ overlap_threshold: float,
+ metadata,
+ size_divisibility: int,
+ sem_seg_postprocess_before_inference: bool,
+ pixel_mean: Tuple[float],
+ pixel_std: Tuple[float],
+ # inference
+ semantic_on: bool,
+ panoptic_on: bool,
+ instance_on: bool,
+ test_topk_per_image: int,
+ ):
+ """
+ Args:
+ backbone: a backbone module, must follow detectron2's backbone interface
+ sem_seg_head: a module that predicts semantic segmentation from backbone features
+ criterion: a module that defines the loss
+ num_queries: int, number of queries
+ object_mask_threshold: float, threshold to filter query based on classification score
+ for panoptic segmentation inference
+ overlap_threshold: overlap threshold used in general inference for panoptic segmentation
+ metadata: dataset meta, get `thing` and `stuff` category names for panoptic
+ segmentation inference
+ size_divisibility: Some backbones require the input height and width to be divisible by a
+ specific integer. We can use this to override such requirement.
+ sem_seg_postprocess_before_inference: whether to resize the prediction back
+ to original input size before semantic segmentation inference or after.
+ For high-resolution dataset like Mapillary, resizing predictions before
+ inference will cause OOM error.
+ pixel_mean, pixel_std: list or tuple with #channels element, representing
+ the per-channel mean and std to be used to normalize the input image
+ semantic_on: bool, whether to output semantic segmentation prediction
+ instance_on: bool, whether to output instance segmentation prediction
+ panoptic_on: bool, whether to output panoptic segmentation prediction
+ test_topk_per_image: int, instance segmentation parameter, keep topk instances per image
+ """
+ super().__init__()
+ self.cfg = cfg
+ self.backbone = backbone
+ self.sem_seg_head = sem_seg_head
+ self.criterion = criterion
+ self.num_queries = num_queries
+ self.overlap_threshold = overlap_threshold
+ self.entity_enable = self.cfg.ENTITY.ENABLE
+ self.object_mask_threshold = object_mask_threshold
+ self.metadata = metadata
+ if size_divisibility < 0:
+ # use backbone size_divisibility if not set
+ size_divisibility = self.backbone.size_divisibility
+ self.size_divisibility = size_divisibility
+ self.sem_seg_postprocess_before_inference = sem_seg_postprocess_before_inference
+ self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False)
+ self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)
+
+ # additional args
+ self.semantic_on = semantic_on
+ self.instance_on = instance_on
+ self.panoptic_on = panoptic_on
+ self.test_topk_per_image = test_topk_per_image
+
+ if not self.semantic_on:
+ assert self.sem_seg_postprocess_before_inference
+
+ @classmethod
+ def from_config(cls, cfg):
+ backbone = build_backbone(cfg)
+ sem_seg_head = build_sem_seg_head(cfg, backbone.output_shape())
+
+ # Loss parameters:
+ deep_supervision = cfg.MODEL.MASK_FORMER.DEEP_SUPERVISION
+ no_object_weight = cfg.MODEL.MASK_FORMER.NO_OBJECT_WEIGHT
+
+ # loss weights
+ class_weight = cfg.MODEL.MASK_FORMER.CLASS_WEIGHT
+ dice_weight = cfg.MODEL.MASK_FORMER.DICE_WEIGHT
+ mask_weight = cfg.MODEL.MASK_FORMER.MASK_WEIGHT
+
+ # building criterion
+ matcher = HungarianMatcher(
+ cost_class=class_weight,
+ cost_mask=mask_weight,
+ cost_dice=dice_weight,
+ num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS,
+ )
+
+ weight_dict = {"loss_ce": class_weight, "loss_mask": mask_weight, "loss_dice": dice_weight}
+
+ if deep_supervision:
+ dec_layers = cfg.MODEL.MASK_FORMER.DEC_LAYERS
+ aux_weight_dict = {}
+ for i in range(dec_layers - 1):
+ aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
+ weight_dict.update(aux_weight_dict)
+
+ losses = ["labels", "masks"]
+
+ criterion = SetCriterion(
+ sem_seg_head.num_classes,
+ matcher=matcher,
+ weight_dict=weight_dict,
+ eos_coef=no_object_weight,
+ losses=losses,
+ num_points=cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS,
+ oversample_ratio=cfg.MODEL.MASK_FORMER.OVERSAMPLE_RATIO,
+ importance_sample_ratio=cfg.MODEL.MASK_FORMER.IMPORTANCE_SAMPLE_RATIO,
+ )
+
+ return {
+ "cfg": cfg,
+ "backbone": backbone,
+ "sem_seg_head": sem_seg_head,
+ "criterion": criterion,
+ "num_queries": cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES,
+ "object_mask_threshold": cfg.MODEL.MASK_FORMER.TEST.OBJECT_MASK_THRESHOLD,
+ "overlap_threshold": cfg.MODEL.MASK_FORMER.TEST.OVERLAP_THRESHOLD,
+ "metadata": MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
+ "size_divisibility": cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY,
+ "sem_seg_postprocess_before_inference": (
+ cfg.MODEL.MASK_FORMER.TEST.SEM_SEG_POSTPROCESSING_BEFORE_INFERENCE
+ or cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON
+ or cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON
+ ),
+ "pixel_mean": cfg.MODEL.PIXEL_MEAN,
+ "pixel_std": cfg.MODEL.PIXEL_STD,
+ # inference
+ "semantic_on": cfg.MODEL.MASK_FORMER.TEST.SEMANTIC_ON,
+ "instance_on": cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON,
+ "panoptic_on": cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON,
+ "test_topk_per_image": cfg.TEST.DETECTIONS_PER_IMAGE,
+ }
+
+ @property
+ def device(self):
+ return self.pixel_mean.device
+
+ def forward(self, batched_inputs):
+ """
+ Args:
+ batched_inputs: a list, batched outputs of :class:`DatasetMapper`.
+ Each item in the list contains the inputs for one image.
+ For now, each item in the list is a dict that contains:
+ * "image": Tensor, image in (C, H, W) format.
+ * "instances": per-region ground truth
+ * Other information that's included in the original dicts, such as:
+ "height", "width" (int): the output resolution of the model (may be different
+ from input resolution), used in inference.
+ Returns:
+ list[dict]:
+ each dict has the results for one image. The dict contains the following keys:
+
+ * "sem_seg":
+ A Tensor that represents the
+ per-pixel segmentation prediced by the head.
+ The prediction has shape KxHxW that represents the logits of
+ each class for each pixel.
+ * "panoptic_seg":
+ A tuple that represent panoptic output
+ panoptic_seg (Tensor): of shape (height, width) where the values are ids for each segment.
+ segments_info (list[dict]): Describe each segment in `panoptic_seg`.
+ Each dict contains keys "id", "category_id", "isthing".
+ """
+ images = [x["image"].to(self.device) for x in batched_inputs]
+ images = [(x - self.pixel_mean) / self.pixel_std for x in images]
+ images = ImageList.from_tensors(images, self.size_divisibility)
+
+ features = self.backbone(images.tensor)
+ outputs = self.sem_seg_head(features)
+
+ if self.training:
+ # mask classification target
+ if "instances" in batched_inputs[0]:
+ if self.cfg.ENTITY.ENABLE:
+ for i in range(len(batched_inputs)):
+ batched_inputs[i]["instances"].gt_classes[:] = 0
+ gt_instances = [x["instances"].to(self.device) for x in batched_inputs]
+ targets = self.prepare_targets(gt_instances, images)
+ else:
+ targets = None
+
+ # bipartite matching-based loss
+ losses = self.criterion(outputs, targets)
+
+ for k in list(losses.keys()):
+ if k in self.criterion.weight_dict:
+ losses[k] *= self.criterion.weight_dict[k]
+ else:
+ # remove this loss if not specified in `weight_dict`
+ losses.pop(k)
+ return losses
+ else:
+ mask_cls_results = outputs["pred_logits"]
+ mask_pred_results = outputs["pred_masks"]
+ # upsample masks
+ mask_pred_results = F.interpolate(
+ mask_pred_results,
+ size=(images.tensor.shape[-2], images.tensor.shape[-1]),
+ mode="bilinear",
+ align_corners=False,
+ )
+
+ del outputs
+
+ processed_results = []
+ for mask_cls_result, mask_pred_result, input_per_image, image_size in zip(
+ mask_cls_results, mask_pred_results, batched_inputs, images.image_sizes
+ ):
+ height = input_per_image.get("height", image_size[0])
+ width = input_per_image.get("width", image_size[1])
+ processed_results.append({})
+
+ if self.sem_seg_postprocess_before_inference:
+ mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
+ mask_pred_result, image_size, height, width
+ )
+ mask_cls_result = mask_cls_result.to(mask_pred_result)
+
+ # semantic segmentation inference
+ if self.semantic_on:
+ r = retry_if_cuda_oom(self.semantic_inference)(mask_cls_result, mask_pred_result)
+ if not self.sem_seg_postprocess_before_inference:
+ r = retry_if_cuda_oom(sem_seg_postprocess)(r, image_size, height, width)
+ processed_results[-1]["sem_seg"] = r
+
+ # panoptic segmentation inference
+ if self.panoptic_on:
+ panoptic_r = retry_if_cuda_oom(self.panoptic_inference)(mask_cls_result, mask_pred_result)
+ processed_results[-1]["panoptic_seg"] = panoptic_r
+ 
+ # instance segmentation and entity segmentation inference
+ if self.instance_on and self.cfg.ENTITY.ENABLE:
+ instance_r = retry_if_cuda_oom(self.instance_inference_nonoverlap)(mask_cls_result, mask_pred_result)
+ processed_results[-1]["instances"] = instance_r
+ else:
+ instance_r = retry_if_cuda_oom(self.instance_inference)(mask_cls_result, mask_pred_result)
+ processed_results[-1]["instances"] = instance_r
+
+ return processed_results
+
+ def prepare_targets(self, targets, images):
+ h_pad, w_pad = images.tensor.shape[-2:]
+ new_targets = []
+ for targets_per_image in targets:
+ # pad gt
+ gt_masks = targets_per_image.gt_masks
+ padded_masks = torch.zeros((gt_masks.shape[0], h_pad, w_pad), dtype=gt_masks.dtype, device=gt_masks.device)
+ padded_masks[:, : gt_masks.shape[1], : gt_masks.shape[2]] = gt_masks
+ new_targets.append(
+ {
+ "labels": targets_per_image.gt_classes,
+ "masks": padded_masks,
+ }
+ )
+ return new_targets
+
+ def semantic_inference(self, mask_cls, mask_pred):
+ mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1]
+ mask_pred = mask_pred.sigmoid()
+ semseg = torch.einsum("qc,qhw->chw", mask_cls, mask_pred)
+ return semseg
+
+ def panoptic_inference(self, mask_cls, mask_pred):
+ scores, labels = F.softmax(mask_cls, dim=-1).max(-1)
+ mask_pred = mask_pred.sigmoid()
+
+ keep = labels.ne(self.sem_seg_head.num_classes) & (scores > self.object_mask_threshold)
+ cur_scores = scores[keep]
+ cur_classes = labels[keep]
+ cur_masks = mask_pred[keep]
+ cur_mask_cls = mask_cls[keep]
+ cur_mask_cls = cur_mask_cls[:, :-1]
+
+ cur_prob_masks = cur_scores.view(-1, 1, 1) * cur_masks
+
+ h, w = cur_masks.shape[-2:]
+ panoptic_seg = torch.zeros((h, w), dtype=torch.int32, device=cur_masks.device)
+ segments_info = []
+
+ current_segment_id = 0
+
+ if cur_masks.shape[0] == 0:
+ # We didn't detect any mask :(
+ return panoptic_seg, segments_info
+ else:
+ # take argmax
+ cur_mask_ids = cur_prob_masks.argmax(0)
+ stuff_memory_list = {}
+ for k in range(cur_classes.shape[0]):
+ pred_class = cur_classes[k].item()
+ isthing = pred_class in self.metadata.thing_dataset_id_to_contiguous_id.values()
+ mask_area = (cur_mask_ids == k).sum().item()
+ original_area = (cur_masks[k] >= 0.5).sum().item()
+ mask = (cur_mask_ids == k) & (cur_masks[k] >= 0.5)
+
+ if mask_area > 0 and original_area > 0 and mask.sum().item() > 0:
+ if mask_area / original_area < self.overlap_threshold:
+ continue
+
+ # merge stuff regions
+ if not isthing:
+ if int(pred_class) in stuff_memory_list.keys():
+ panoptic_seg[mask] = stuff_memory_list[int(pred_class)]
+ continue
+ else:
+ stuff_memory_list[int(pred_class)] = current_segment_id + 1
+
+ current_segment_id += 1
+ panoptic_seg[mask] = current_segment_id
+
+ segments_info.append(
+ {
+ "id": current_segment_id,
+ "isthing": bool(isthing),
+ "category_id": int(pred_class),
+ }
+ )
+
+ return panoptic_seg, segments_info
+
+ def instance_inference(self, mask_cls, mask_pred):
+ # mask_pred is already processed to have the same shape as original input
+ image_size = mask_pred.shape[-2:]
+
+ # [Q, K]
+ scores = F.softmax(mask_cls, dim=-1)[:, :-1]
+ labels = torch.arange(self.sem_seg_head.num_classes, device=self.device).unsqueeze(0).repeat(self.num_queries, 1).flatten(0, 1)
+ # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
+ scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.test_topk_per_image, sorted=False)
+ labels_per_image = labels[topk_indices]
+
+ # topk_indices = topk_indices // self.sem_seg_head.num_classes
+ topk_indices = torch.div(topk_indices, self.sem_seg_head.num_classes, rounding_mode='trunc')
+ # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
+ mask_pred = mask_pred[topk_indices]
+
+ # if this is panoptic segmentation, we only keep the "thing" classes
+ if self.panoptic_on:
+ keep = torch.zeros_like(scores_per_image).bool()
+ for i, lab in enumerate(labels_per_image):
+ keep[i] = lab in self.metadata.thing_dataset_id_to_contiguous_id.values()
+
+ scores_per_image = scores_per_image[keep]
+ labels_per_image = labels_per_image[keep]
+ mask_pred = mask_pred[keep]
+
+ result = Instances(image_size)
+ # mask (before sigmoid)
+ result.pred_masks = (mask_pred > 0).float()
+ result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
+ # Uncomment the following to get boxes from masks (this is slow)
+ # result.pred_boxes = BitMasks(mask_pred > 0).get_bounding_boxes()
+
+ # calculate average mask prob
+ mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
+ result.scores = scores_per_image * mask_scores_per_image
+ result.pred_classes = labels_per_image
+ return result
+ 
+ def instance_inference_nonoverlap(self, mask_cls, mask_pred):
+ # mask_pred is already processed to have the same shape as original input
+ image_size = mask_pred.shape[-2:]
+
+ # [Q, K]
+ scores = F.softmax(mask_cls, dim=-1)[:, :-1]
+ labels = torch.arange(self.sem_seg_head.num_classes, device=self.device).unsqueeze(0).repeat(self.num_queries, 1).flatten(0, 1)
+ # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
+ scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.test_topk_per_image, sorted=False)
+ labels_per_image = labels[topk_indices]
+
+ # topk_indices = topk_indices // self.sem_seg_head.num_classes
+ topk_indices = torch.div(topk_indices, self.sem_seg_head.num_classes, rounding_mode='trunc')
+ # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
+ mask_pred = mask_pred[topk_indices]
+
+ ###### ranks
+ pred_masks = (mask_pred>0).float()
+ pred_masks_logits = mask_pred.sigmoid()
+ pred_scores = scores_per_image
+
+ _, m_H, m_W = pred_masks.shape
+ mask_id = torch.zeros((m_H, m_W), dtype=torch.int).to(pred_masks.device)
+ sorted_scores, ranks = torch.sort(pred_scores)
+ ranks = ranks + 1
+ for index in ranks:
+ mask_id[(pred_masks[index-1]==1)] = int(index)
+ # re-generate mask
+ new_scores = []
+ new_masks = []
+ new_masks_logits = []
+ entity_nums = len(ranks)
+ for ii in range(entity_nums):
+ index = int(ranks[entity_nums-ii-1])
+ score = sorted_scores[entity_nums-ii-1]
+ new_scores.append(score)
+ new_masks.append((mask_id==index).float())
+ new_masks_logits.append(pred_masks_logits[index-1])
+ 
+ new_scores = torch.stack(new_scores)
+ new_masks = torch.stack(new_masks)
+ new_masks_logits = torch.stack(new_masks_logits)
+
+ result = Instances(image_size)
+ # mask (before sigmoid)
+ result.pred_masks = new_masks
+ result.pred_boxes = Boxes(torch.zeros(new_masks.size(0), 4))
+ # Uncomment the following to get boxes from masks (this is slow)
+
+ # calculate average mask prob
+ mask_scores_per_image = (new_masks_logits.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
+ result.scores = new_scores * mask_scores_per_image
+ result.pred_classes = labels_per_image
+ return result

annotator/entityseg/mask2former/modeling/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from .backbone.swin import D2SwinTransformer
+from .backbone.hornet import D2HorNet
+from .pixel_decoder.fpn import BasePixelDecoder
+from .pixel_decoder.msdeformattn import MSDeformAttnPixelDecoder
+from .meta_arch.mask_former_head import MaskFormerHead
+from .meta_arch.per_pixel_baseline import PerPixelBaselineHead, PerPixelBaselinePlusHead

annotator/entityseg/mask2former/modeling/backbone/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Facebook, Inc. and its affiliates.

annotator/entityseg/mask2former/modeling/backbone/hornet.py

@@ -0,0 +1,363 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from functools import partial
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import trunc_normal_, DropPath
+from timm.models.registry import register_model
+import os
+import sys
+import torch.fft
+import math
+
+import traceback
+import torch.utils.checkpoint as checkpoint
+from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec
+
+
+if 'DWCONV_IMPL' in os.environ:
+ try:
+ sys.path.append(os.environ['DWCONV_IMPL'])
+ from depthwise_conv2d_implicit_gemm import DepthWiseConv2dImplicitGEMM
+ def get_dwconv(dim, kernel, bias):
+ return DepthWiseConv2dImplicitGEMM(dim, kernel, bias)
+ print('Using Megvii large kernel dw conv impl')
+ except:
+ print(traceback.format_exc())
+ def get_dwconv(dim, kernel, bias):
+ return nn.Conv2d(dim, dim, kernel_size=kernel, padding=(kernel-1)//2 ,bias=bias, groups=dim)
+
+ print('[fail to use Megvii Large kernel] Using PyTorch large kernel dw conv impl')
+else:
+ def get_dwconv(dim, kernel, bias):
+ return nn.Conv2d(dim, dim, kernel_size=kernel, padding=(kernel-1)//2 ,bias=bias, groups=dim)
+
+ print('Using PyTorch large kernel dw conv impl')
+
+class GlobalLocalFilter(nn.Module):
+ def __init__(self, dim, h=14, w=8):
+ super().__init__()
+ self.dw = nn.Conv2d(dim // 2, dim // 2, kernel_size=3, padding=1, bias=False, groups=dim // 2)
+ self.complex_weight = nn.Parameter(torch.randn(dim // 2, h, w, 2, dtype=torch.float32) * 0.02)
+ trunc_normal_(self.complex_weight, std=.02)
+ self.pre_norm = LayerNorm(dim, eps=1e-6, data_format='channels_first')
+ self.post_norm = LayerNorm(dim, eps=1e-6, data_format='channels_first')
+
+ def forward(self, x):
+ x = self.pre_norm(x)
+ x1, x2 = torch.chunk(x, 2, dim=1)
+ x1 = self.dw(x1)
+
+ x2 = x2.to(torch.float32)
+ B, C, a, b = x2.shape
+ x2 = torch.fft.rfft2(x2, dim=(2, 3), norm='ortho')
+
+ weight = self.complex_weight
+ if not weight.shape[1:3] == x2.shape[2:4]:
+ weight = F.interpolate(weight.permute(3,0,1,2), size=x2.shape[2:4], mode='bilinear', align_corners=True).permute(1,2,3,0)
+
+ weight = torch.view_as_complex(weight.contiguous())
+
+ x2 = x2 * weight
+ x2 = torch.fft.irfft2(x2, s=(a, b), dim=(2, 3), norm='ortho')
+
+ x = torch.cat([x1.unsqueeze(2), x2.unsqueeze(2)], dim=2).reshape(B, 2 * C, a, b)
+ x = self.post_norm(x)
+ return x
+
+
+class gnconv(nn.Module):
+ def __init__(self, dim, order=5, gflayer=None, h=14, w=8, s=1.0):
+ super().__init__()
+ self.order = order
+ self.dims = [dim // 2 ** i for i in range(order)]
+ self.dims.reverse()
+ self.proj_in = nn.Conv2d(dim, 2*dim, 1)
+
+ if gflayer is None:
+ self.dwconv = get_dwconv(sum(self.dims), 7, True)
+ else:
+ self.dwconv = gflayer(sum(self.dims), h=h, w=w)
+ 
+ self.proj_out = nn.Conv2d(dim, dim, 1)
+
+ self.pws = nn.ModuleList(
+ [nn.Conv2d(self.dims[i], self.dims[i+1], 1) for i in range(order-1)]
+ )
+
+ self.scale = s
+
+ print('[gconv]', order, 'order with dims=', self.dims, 'scale=%.4f'%self.scale)
+
+
+ def forward(self, x, mask=None, dummy=False):
+ B, C, H, W = x.shape
+
+ fused_x = self.proj_in(x)
+ pwa, abc = torch.split(fused_x, (self.dims[0], sum(self.dims)), dim=1)
+
+ dw_abc = self.dwconv(abc) * self.scale
+
+ dw_list = torch.split(dw_abc, self.dims, dim=1)
+ x = pwa * dw_list[0]
+
+ for i in range(self.order -1):
+ x = self.pws[i](x) * dw_list[i+1]
+
+ x = self.proj_out(x)
+
+ return x
+
+class Block(nn.Module):
+ r""" HorNet block
+ """
+ def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, gnconv=gnconv):
+ super().__init__()
+
+ self.norm1 = LayerNorm(dim, eps=1e-6, data_format='channels_first')
+ self.gnconv = gnconv(dim) # depthwise conv
+ self.norm2 = LayerNorm(dim, eps=1e-6)
+ self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
+ self.act = nn.GELU()
+ self.pwconv2 = nn.Linear(4 * dim, dim)
+
+ self.gamma1 = nn.Parameter(layer_scale_init_value * torch.ones(dim), 
+ requires_grad=True) if layer_scale_init_value > 0 else None
+
+ self.gamma2 = nn.Parameter(layer_scale_init_value * torch.ones((dim)), 
+ requires_grad=True) if layer_scale_init_value > 0 else None
+ self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+ def forward(self, x):
+ B, C, H, W = x.shape
+ if self.gamma1 is not None:
+ gamma1 = self.gamma1.view(C, 1, 1)
+ else:
+ gamma1 = 1
+ x = x + self.drop_path(gamma1 * self.gnconv(self.norm1(x)))
+
+ input = x
+ x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+ x = self.norm2(x)
+ x = self.pwconv1(x)
+ x = self.act(x)
+ x = self.pwconv2(x)
+ if self.gamma2 is not None:
+ x = self.gamma2 * x
+ x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
+
+ x = input + self.drop_path(x)
+ return x
+
+
+class HorNet(nn.Module):
+ r""" HorNet
+ A PyTorch impl of : `HorNet: Efficient High-Order Spatial Interactions with Recursive Gated Convolutions`
+
+ Args:
+ in_chans (int): Number of input image channels. Default: 3
+ num_classes (int): Number of classes for classification head. Default: 1000
+ depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+ dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+ drop_path_rate (float): Stochastic depth rate. Default: 0.
+ layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+ head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+ """
+ def __init__(self, in_chans=3, num_classes=1000, 
+ depths=[3, 3, 9, 3], base_dim=96, drop_path_rate=0.,
+ layer_scale_init_value=1e-6, head_init_scale=1.,
+ gnconv=gnconv, block=Block,
+ pretrained=None,
+ use_checkpoint=False,
+ ):
+ super().__init__()
+
+ self.pretrained = pretrained
+ self.use_checkpoint = use_checkpoint
+
+ dims = [base_dim, base_dim*2, base_dim*4, base_dim*8]
+
+ self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
+ stem = nn.Sequential(
+ nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+ LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
+ )
+ self.downsample_layers.append(stem)
+ for i in range(3):
+ downsample_layer = nn.Sequential(
+ LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+ nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
+ )
+ self.downsample_layers.append(downsample_layer)
+
+ self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
+ dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] 
+
+
+ if not isinstance(gnconv, list):
+ gnconv = [gnconv, gnconv, gnconv, gnconv]
+ else:
+ gnconv = gnconv
+ assert len(gnconv) == 4
+
+ if isinstance(gnconv[0], str):
+ print('[GConvNet]: convert str gconv to func')
+ gnconv = [eval(g) for g in gnconv]
+
+ if isinstance(block, str):
+ block = eval(block)
+
+ cur = 0
+ num_features = []
+ for i in range(4):
+ stage = nn.Sequential(
+ *[block(dim=dims[i], drop_path=dp_rates[cur + j], 
+ layer_scale_init_value=layer_scale_init_value, gnconv=gnconv[i]) for j in range(depths[i])]
+ )
+ self.stages.append(stage)
+ cur += depths[i]
+ num_features.append(dims[i])
+ self.num_features = num_features
+
+ norm_layer = partial(LayerNorm, eps=1e-6, data_format="channels_first")
+ for i_layer in range(4):
+ layer = norm_layer(dims[i_layer])
+ layer_name = f'norm{i_layer}'
+ self.add_module(layer_name, layer)
+
+ def init_weights(self):
+ """Initialize the weights in backbone.
+ Args:
+ pretrained (str, optional): Path to pre-trained weights.
+ Defaults to None.
+ """
+ #pretrained = self.pretrained
+
+ def _init_weights(m):
+ if isinstance(m, nn.Linear):
+ trunc_normal_(m.weight, std=.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)
+
+ #if isinstance(pretrained, str):
+ # self.apply(_init_weights)
+ # logger = get_root_logger()
+ # load_checkpoint(self, pretrained, strict=False, logger=logger)
+ #elif pretrained is None:
+ # raise NotImplementedError()
+ self.apply(_init_weights)
+ #else:
+ # raise TypeError('pretrained must be a str or None')
+
+ def forward_features(self, x):
+ outs = dict()
+ for i in range(4):
+ x = self.downsample_layers[i](x)
+ if self.use_checkpoint:
+ x = checkpoint.checkpoint_sequential(self.stages[i], len(self.stages[i]), x)
+ else:
+ x = self.stages[i](x)
+ norm_layer = getattr(self, f'norm{i}')
+ x_out = norm_layer(x)
+ outs["res%i"% (i+2)] = x_out
+ return outs #tuple(outs)
+
+ def forward(self, x):
+ x = self.forward_features(x)
+ return x
+
+
+class LayerNorm(nn.Module):
+ r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
+ The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
+ shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
+ with shape (batch_size, channels, height, width).
+ """
+ def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(normalized_shape))
+ self.bias = nn.Parameter(torch.zeros(normalized_shape))
+ self.eps = eps
+ self.data_format = data_format
+ if self.data_format not in ["channels_last", "channels_first"]:
+ raise NotImplementedError 
+ self.normalized_shape = (normalized_shape, )
+ 
+ def forward(self, x):
+ if self.data_format == "channels_last":
+ return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+ elif self.data_format == "channels_first":
+ u = x.mean(1, keepdim=True)
+ s = (x - u).pow(2).mean(1, keepdim=True)
+ x = (x - u) / torch.sqrt(s + self.eps)
+ x = self.weight[:, None, None] * x + self.bias[:, None, None]
+ return x
+ 
+@BACKBONE_REGISTRY.register()
+class D2HorNet(HorNet, Backbone):
+ def __init__(self, cfg, input_shape):
+
+ depths=cfg.MODEL.HORNET.DEPTHS
+ base_dim=cfg.MODEL.HORNET.BASE_DIM
+ gnconv=cfg.MODEL.HORNET.GCONV
+ drop_path_rate=cfg.MODEL.HORNET.DROP_PATH_RATE
+
+ super().__init__(
+ depths=depths,
+ base_dim=base_dim,
+ gnconv=gnconv,
+ drop_path_rate=drop_path_rate,
+ )
+
+ self._out_features = cfg.MODEL.HORNET.OUT_FEATURES
+
+ self._out_feature_strides = {
+ "res2": 4,
+ "res3": 8,
+ "res4": 16,
+ "res5": 32,
+ }
+ self._out_feature_channels = {
+ "res2": self.num_features[0],
+ "res3": self.num_features[1],
+ "res4": self.num_features[2],
+ "res5": self.num_features[3],
+ }
+
+ def forward(self, x):
+ """
+ Args:
+ x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
+ Returns:
+ dict[str->Tensor]: names and the corresponding features
+ """
+ assert (
+ x.dim() == 4
+ ), f"SwinTransformer takes an input of shape (N, C, H, W). Got {x.shape} instead!"
+ outputs = {}
+ y = super().forward(x)
+ for k in y.keys():
+ if k in self._out_features:
+ outputs[k] = y[k]
+ return outputs
+
+ 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
+ }
+
+ @property
+ def size_divisibility(self):
+ return 32

annotator/entityseg/mask2former/modeling/backbone/swin.py

@@ -0,0 +1,770 @@
+# --------------------------------------------------------
+# Swin Transformer
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ze Liu, Yutong Lin, Yixuan Wei
+# --------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation/blob/main/mmseg/models/backbones/swin_transformer.py
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec
+
+
+class Mlp(nn.Module):
+ """Multilayer perceptron."""
+
+ def __init__(
+ self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0
+ ):
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ self.fc1 = nn.Linear(in_features, hidden_features)
+ self.act = act_layer()
+ self.fc2 = nn.Linear(hidden_features, out_features)
+ self.drop = nn.Dropout(drop)
+
+ def forward(self, x):
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.drop(x)
+ x = self.fc2(x)
+ x = self.drop(x)
+ return x
+
+
+def window_partition(x, window_size):
+ """
+ Args:
+ x: (B, H, W, C)
+ window_size (int): window size
+ Returns:
+ windows: (num_windows*B, window_size, window_size, C)
+ """
+ B, H, W, C = x.shape
+ x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+ return windows
+
+
+def window_reverse(windows, window_size, H, W):
+ """
+ Args:
+ windows: (num_windows*B, window_size, window_size, C)
+ window_size (int): Window size
+ H (int): Height of image
+ W (int): Width of image
+ Returns:
+ x: (B, H, W, C)
+ """
+ B = int(windows.shape[0] / (H * W / window_size / window_size))
+ x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+ return x
+
+
+class WindowAttention(nn.Module):
+ """Window based multi-head self attention (W-MSA) module with relative position bias.
+ It supports both of shifted and non-shifted window.
+ Args:
+ dim (int): Number of input channels.
+ window_size (tuple[int]): The height and width of the window.
+ num_heads (int): Number of attention heads.
+ qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+ qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+ attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+ proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+ """
+
+ def __init__(
+ self,
+ dim,
+ window_size,
+ num_heads,
+ qkv_bias=True,
+ qk_scale=None,
+ attn_drop=0.0,
+ proj_drop=0.0,
+ ):
+
+ super().__init__()
+ self.dim = dim
+ self.window_size = window_size # Wh, Ww
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ self.scale = qk_scale or head_dim ** -0.5
+
+ # define a parameter table of relative position bias
+ self.relative_position_bias_table = nn.Parameter(
+ torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
+ ) # 2*Wh-1 * 2*Ww-1, nH
+
+ # get pair-wise relative position index for each token inside the window
+ coords_h = torch.arange(self.window_size[0])
+ coords_w = torch.arange(self.window_size[1])
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
+ relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
+ relative_coords[:, :, 1] += self.window_size[1] - 1
+ relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+ relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
+ self.register_buffer("relative_position_index", relative_position_index)
+
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+ self.attn_drop = nn.Dropout(attn_drop)
+ self.proj = nn.Linear(dim, dim)
+ self.proj_drop = nn.Dropout(proj_drop)
+
+ trunc_normal_(self.relative_position_bias_table, std=0.02)
+ self.softmax = nn.Softmax(dim=-1)
+
+ def forward(self, x, mask=None):
+ """Forward function.
+ Args:
+ x: input features with shape of (num_windows*B, N, C)
+ mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+ """
+ B_, N, C = x.shape
+ qkv = (
+ self.qkv(x)
+ .reshape(B_, N, 3, self.num_heads, C // self.num_heads)
+ .permute(2, 0, 3, 1, 4)
+ )
+ q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
+
+ q = q * self.scale
+ attn = q @ k.transpose(-2, -1)
+
+ relative_position_bias = self.relative_position_bias_table[
+ self.relative_position_index.view(-1)
+ ].view(
+ self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
+ ) # Wh*Ww,Wh*Ww,nH
+ relative_position_bias = relative_position_bias.permute(
+ 2, 0, 1
+ ).contiguous() # nH, Wh*Ww, Wh*Ww
+ attn = attn + relative_position_bias.unsqueeze(0)
+
+ if mask is not None:
+ nW = mask.shape[0]
+ attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+ attn = attn.view(-1, self.num_heads, N, N)
+ attn = self.softmax(attn)
+ else:
+ attn = self.softmax(attn)
+
+ attn = self.attn_drop(attn)
+
+ x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
+
+
+class SwinTransformerBlock(nn.Module):
+ """Swin Transformer Block.
+ Args:
+ dim (int): Number of input channels.
+ num_heads (int): Number of attention heads.
+ window_size (int): Window size.
+ shift_size (int): Shift size for SW-MSA.
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+ qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+ qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+ drop (float, optional): Dropout rate. Default: 0.0
+ attn_drop (float, optional): Attention dropout rate. Default: 0.0
+ drop_path (float, optional): Stochastic depth rate. Default: 0.0
+ act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+ norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+ """
+
+ def __init__(
+ self,
+ dim,
+ num_heads,
+ window_size=7,
+ shift_size=0,
+ mlp_ratio=4.0,
+ qkv_bias=True,
+ qk_scale=None,
+ drop=0.0,
+ attn_drop=0.0,
+ drop_path=0.0,
+ act_layer=nn.GELU,
+ norm_layer=nn.LayerNorm,
+ ):
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.window_size = window_size
+ self.shift_size = shift_size
+ self.mlp_ratio = mlp_ratio
+ assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+ self.norm1 = norm_layer(dim)
+ self.attn = WindowAttention(
+ dim,
+ window_size=to_2tuple(self.window_size),
+ num_heads=num_heads,
+ qkv_bias=qkv_bias,
+ qk_scale=qk_scale,
+ attn_drop=attn_drop,
+ proj_drop=drop,
+ )
+
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+ self.norm2 = norm_layer(dim)
+ mlp_hidden_dim = int(dim * mlp_ratio)
+ self.mlp = Mlp(
+ in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop
+ )
+
+ self.H = None
+ self.W = None
+
+ def forward(self, x, mask_matrix):
+ """Forward function.
+ Args:
+ x: Input feature, tensor size (B, H*W, C).
+ H, W: Spatial resolution of the input feature.
+ mask_matrix: Attention mask for cyclic shift.
+ """
+ B, L, C = x.shape
+ H, W = self.H, self.W
+ assert L == H * W, "input feature has wrong size"
+
+ shortcut = x
+ x = self.norm1(x)
+ x = x.view(B, H, W, C)
+
+ # pad feature maps to multiples of window size
+ pad_l = pad_t = 0
+ pad_r = (self.window_size - W % self.window_size) % self.window_size
+ pad_b = (self.window_size - H % self.window_size) % self.window_size
+ x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
+ _, Hp, Wp, _ = x.shape
+
+ # cyclic shift
+ if self.shift_size > 0:
+ shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+ attn_mask = mask_matrix
+ else:
+ shifted_x = x
+ attn_mask = None
+
+ # partition windows
+ x_windows = window_partition(
+ shifted_x, self.window_size
+ ) # nW*B, window_size, window_size, C
+ x_windows = x_windows.view(
+ -1, self.window_size * self.window_size, C
+ ) # nW*B, window_size*window_size, C
+
+ # W-MSA/SW-MSA
+ attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C
+
+ # merge windows
+ attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+ shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C
+
+ # reverse cyclic shift
+ if self.shift_size > 0:
+ x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+ else:
+ x = shifted_x
+
+ if pad_r > 0 or pad_b > 0:
+ x = x[:, :H, :W, :].contiguous()
+
+ x = x.view(B, H * W, C)
+
+ # FFN
+ x = shortcut + self.drop_path(x)
+ x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+ return x
+
+
+class PatchMerging(nn.Module):
+ """Patch Merging Layer
+ Args:
+ dim (int): Number of input channels.
+ norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+ """
+
+ def __init__(self, dim, norm_layer=nn.LayerNorm):
+ super().__init__()
+ self.dim = dim
+ self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+ self.norm = norm_layer(4 * dim)
+
+ def forward(self, x, H, W):
+ """Forward function.
+ Args:
+ x: Input feature, tensor size (B, H*W, C).
+ H, W: Spatial resolution of the input feature.
+ """
+ B, L, C = x.shape
+ assert L == H * W, "input feature has wrong size"
+
+ x = x.view(B, H, W, C)
+
+ # padding
+ pad_input = (H % 2 == 1) or (W % 2 == 1)
+ if pad_input:
+ x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
+
+ x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C
+ x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C
+ x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C
+ x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C
+ x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C
+ x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C
+
+ x = self.norm(x)
+ x = self.reduction(x)
+
+ return x
+
+
+class BasicLayer(nn.Module):
+ """A basic Swin Transformer layer for one stage.
+ Args:
+ dim (int): Number of feature channels
+ depth (int): Depths of this stage.
+ num_heads (int): Number of attention head.
+ window_size (int): Local window size. Default: 7.
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+ qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+ qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+ drop (float, optional): Dropout rate. Default: 0.0
+ attn_drop (float, optional): Attention dropout rate. Default: 0.0
+ drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+ norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+ downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+ use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+ """
+
+ def __init__(
+ self,
+ dim,
+ depth,
+ num_heads,
+ window_size=7,
+ mlp_ratio=4.0,
+ qkv_bias=True,
+ qk_scale=None,
+ drop=0.0,
+ attn_drop=0.0,
+ drop_path=0.0,
+ norm_layer=nn.LayerNorm,
+ downsample=None,
+ use_checkpoint=False,
+ ):
+ super().__init__()
+ self.window_size = window_size
+ self.shift_size = window_size // 2
+ self.depth = depth
+ self.use_checkpoint = use_checkpoint
+
+ # build blocks
+ self.blocks = nn.ModuleList(
+ [
+ SwinTransformerBlock(
+ dim=dim,
+ num_heads=num_heads,
+ window_size=window_size,
+ shift_size=0 if (i % 2 == 0) else window_size // 2,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ qk_scale=qk_scale,
+ drop=drop,
+ attn_drop=attn_drop,
+ drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+ norm_layer=norm_layer,
+ )
+ for i in range(depth)
+ ]
+ )
+
+ # patch merging layer
+ if downsample is not None:
+ self.downsample = downsample(dim=dim, norm_layer=norm_layer)
+ else:
+ self.downsample = None
+
+ def forward(self, x, H, W):
+ """Forward function.
+ Args:
+ x: Input feature, tensor size (B, H*W, C).
+ H, W: Spatial resolution of the input feature.
+ """
+
+ # calculate attention mask for SW-MSA
+ Hp = int(np.ceil(H / self.window_size)) * self.window_size
+ Wp = int(np.ceil(W / self.window_size)) * self.window_size
+ img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1
+ h_slices = (
+ slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None),
+ )
+ w_slices = (
+ slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None),
+ )
+ cnt = 0
+ for h in h_slices:
+ for w in w_slices:
+ img_mask[:, h, w, :] = cnt
+ cnt += 1
+
+ mask_windows = window_partition(
+ img_mask, self.window_size
+ ) # nW, window_size, window_size, 1
+ mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
+ attn_mask == 0, float(0.0)
+ )
+
+ for blk in self.blocks:
+ blk.H, blk.W = H, W
+ if self.use_checkpoint:
+ x = checkpoint.checkpoint(blk, x, attn_mask)
+ else:
+ x = blk(x, attn_mask)
+ if self.downsample is not None:
+ x_down = self.downsample(x, H, W)
+ Wh, Ww = (H + 1) // 2, (W + 1) // 2
+ return x, H, W, x_down, Wh, Ww
+ else:
+ return x, H, W, x, H, W
+
+
+class PatchEmbed(nn.Module):
+ """Image to Patch Embedding
+ Args:
+ patch_size (int): Patch token size. Default: 4.
+ in_chans (int): Number of input image channels. Default: 3.
+ embed_dim (int): Number of linear projection output channels. Default: 96.
+ norm_layer (nn.Module, optional): Normalization layer. Default: None
+ """
+
+ def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+ super().__init__()
+ patch_size = to_2tuple(patch_size)
+ self.patch_size = patch_size
+
+ self.in_chans = in_chans
+ self.embed_dim = embed_dim
+
+ self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+ if norm_layer is not None:
+ self.norm = norm_layer(embed_dim)
+ else:
+ self.norm = None
+
+ def forward(self, x):
+ """Forward function."""
+ # padding
+ _, _, H, W = x.size()
+ if W % self.patch_size[1] != 0:
+ x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
+ if H % self.patch_size[0] != 0:
+ x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
+
+ x = self.proj(x) # B C Wh Ww
+ if self.norm is not None:
+ Wh, Ww = x.size(2), x.size(3)
+ x = x.flatten(2).transpose(1, 2)
+ x = self.norm(x)
+ x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
+
+ return x
+
+
+class SwinTransformer(nn.Module):
+ """Swin Transformer backbone.
+ A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` -
+ https://arxiv.org/pdf/2103.14030
+ Args:
+ pretrain_img_size (int): Input image size for training the pretrained model,
+ used in absolute postion embedding. Default 224.
+ patch_size (int | tuple(int)): Patch size. Default: 4.
+ in_chans (int): Number of input image channels. Default: 3.
+ embed_dim (int): Number of linear projection output channels. Default: 96.
+ depths (tuple[int]): Depths of each Swin Transformer stage.
+ num_heads (tuple[int]): Number of attention head of each stage.
+ window_size (int): Window size. Default: 7.
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+ qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+ qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
+ drop_rate (float): Dropout rate.
+ attn_drop_rate (float): Attention dropout rate. Default: 0.
+ drop_path_rate (float): Stochastic depth rate. Default: 0.2.
+ norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+ ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.
+ patch_norm (bool): If True, add normalization after patch embedding. Default: True.
+ out_indices (Sequence[int]): Output from which stages.
+ frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
+ -1 means not freezing any parameters.
+ use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+ """
+
+ def __init__(
+ self,
+ pretrain_img_size=224,
+ patch_size=4,
+ in_chans=3,
+ embed_dim=96,
+ depths=[2, 2, 6, 2],
+ num_heads=[3, 6, 12, 24],
+ window_size=7,
+ mlp_ratio=4.0,
+ qkv_bias=True,
+ qk_scale=None,
+ drop_rate=0.0,
+ attn_drop_rate=0.0,
+ drop_path_rate=0.2,
+ norm_layer=nn.LayerNorm,
+ ape=False,
+ patch_norm=True,
+ out_indices=(0, 1, 2, 3),
+ frozen_stages=-1,
+ use_checkpoint=False,
+ ):
+ super().__init__()
+
+ self.pretrain_img_size = pretrain_img_size
+ self.num_layers = len(depths)
+ self.embed_dim = embed_dim
+ self.ape = ape
+ self.patch_norm = patch_norm
+ self.out_indices = out_indices
+ self.frozen_stages = frozen_stages
+
+ # split image into non-overlapping patches
+ self.patch_embed = PatchEmbed(
+ patch_size=patch_size,
+ in_chans=in_chans,
+ embed_dim=embed_dim,
+ norm_layer=norm_layer if self.patch_norm else None,
+ )
+
+ # absolute position embedding
+ if self.ape:
+ pretrain_img_size = to_2tuple(pretrain_img_size)
+ patch_size = to_2tuple(patch_size)
+ patches_resolution = [
+ pretrain_img_size[0] // patch_size[0],
+ pretrain_img_size[1] // patch_size[1],
+ ]
+
+ self.absolute_pos_embed = nn.Parameter(
+ torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])
+ )
+ trunc_normal_(self.absolute_pos_embed, std=0.02)
+
+ self.pos_drop = nn.Dropout(p=drop_rate)
+
+ # stochastic depth
+ dpr = [
+ x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
+ ] # stochastic depth decay rule
+
+ # build layers
+ self.layers = nn.ModuleList()
+ for i_layer in range(self.num_layers):
+ layer = BasicLayer(
+ dim=int(embed_dim * 2 ** i_layer),
+ depth=depths[i_layer],
+ num_heads=num_heads[i_layer],
+ window_size=window_size,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ qk_scale=qk_scale,
+ drop=drop_rate,
+ attn_drop=attn_drop_rate,
+ drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],
+ norm_layer=norm_layer,
+ downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+ use_checkpoint=use_checkpoint,
+ )
+ self.layers.append(layer)
+
+ num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
+ self.num_features = num_features
+
+ # add a norm layer for each output
+ for i_layer in out_indices:
+ layer = norm_layer(num_features[i_layer])
+ layer_name = f"norm{i_layer}"
+ self.add_module(layer_name, layer)
+
+ self._freeze_stages()
+
+ def _freeze_stages(self):
+ if self.frozen_stages >= 0:
+ self.patch_embed.eval()
+ for param in self.patch_embed.parameters():
+ param.requires_grad = False
+
+ if self.frozen_stages >= 1 and self.ape:
+ self.absolute_pos_embed.requires_grad = False
+
+ if self.frozen_stages >= 2:
+ self.pos_drop.eval()
+ for i in range(0, self.frozen_stages - 1):
+ m = self.layers[i]
+ m.eval()
+ for param in m.parameters():
+ param.requires_grad = False
+
+ def init_weights(self, pretrained=None):
+ """Initialize the weights in backbone.
+ Args:
+ pretrained (str, optional): Path to pre-trained weights.
+ Defaults to None.
+ """
+
+ def _init_weights(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):
+ """Forward function."""
+ x = self.patch_embed(x)
+
+ Wh, Ww = x.size(2), x.size(3)
+ if self.ape:
+ # interpolate the position embedding to the corresponding size
+ absolute_pos_embed = F.interpolate(
+ self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic"
+ )
+ x = (x + absolute_pos_embed).flatten(2).transpose(1, 2) # B Wh*Ww C
+ else:
+ x = x.flatten(2).transpose(1, 2)
+ x = self.pos_drop(x)
+
+ outs = {}
+ for i in range(self.num_layers):
+ layer = self.layers[i]
+ x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
+
+ if i in self.out_indices:
+ norm_layer = getattr(self, f"norm{i}")
+ x_out = norm_layer(x_out)
+
+ out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
+ outs["res{}".format(i + 2)] = out
+
+ return outs
+
+ def train(self, mode=True):
+ """Convert the model into training mode while keep layers freezed."""
+ super(SwinTransformer, self).train(mode)
+ self._freeze_stages()
+
+
+@BACKBONE_REGISTRY.register()
+class D2SwinTransformer(SwinTransformer, Backbone):
+ def __init__(self, cfg, input_shape):
+
+ pretrain_img_size = cfg.MODEL.SWIN.PRETRAIN_IMG_SIZE
+ patch_size = cfg.MODEL.SWIN.PATCH_SIZE
+ in_chans = 3
+ embed_dim = cfg.MODEL.SWIN.EMBED_DIM
+ depths = cfg.MODEL.SWIN.DEPTHS
+ num_heads = cfg.MODEL.SWIN.NUM_HEADS
+ window_size = cfg.MODEL.SWIN.WINDOW_SIZE
+ mlp_ratio = cfg.MODEL.SWIN.MLP_RATIO
+ qkv_bias = cfg.MODEL.SWIN.QKV_BIAS
+ qk_scale = cfg.MODEL.SWIN.QK_SCALE
+ drop_rate = cfg.MODEL.SWIN.DROP_RATE
+ attn_drop_rate = cfg.MODEL.SWIN.ATTN_DROP_RATE
+ drop_path_rate = cfg.MODEL.SWIN.DROP_PATH_RATE
+ norm_layer = nn.LayerNorm
+ ape = cfg.MODEL.SWIN.APE
+ patch_norm = cfg.MODEL.SWIN.PATCH_NORM
+ use_checkpoint = cfg.MODEL.SWIN.USE_CHECKPOINT
+
+ super().__init__(
+ pretrain_img_size,
+ patch_size,
+ in_chans,
+ embed_dim,
+ depths,
+ num_heads,
+ window_size,
+ mlp_ratio,
+ qkv_bias,
+ qk_scale,
+ drop_rate,
+ attn_drop_rate,
+ drop_path_rate,
+ norm_layer,
+ ape,
+ patch_norm,
+ use_checkpoint=use_checkpoint,
+ )
+
+ self._out_features = cfg.MODEL.SWIN.OUT_FEATURES
+
+ self._out_feature_strides = {
+ "res2": 4,
+ "res3": 8,
+ "res4": 16,
+ "res5": 32,
+ }
+ self._out_feature_channels = {
+ "res2": self.num_features[0],
+ "res3": self.num_features[1],
+ "res4": self.num_features[2],
+ "res5": self.num_features[3],
+ }
+
+ def forward(self, x):
+ """
+ Args:
+ x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
+ Returns:
+ dict[str->Tensor]: names and the corresponding features
+ """
+ assert (
+ x.dim() == 4
+ ), f"SwinTransformer takes an input of shape (N, C, H, W). Got {x.shape} instead!"
+ outputs = {}
+ y = super().forward(x)
+ for k in y.keys():
+ if k in self._out_features:
+ outputs[k] = y[k]
+ return outputs
+
+ 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
+ }
+
+ @property
+ def size_divisibility(self):
+ return 32

annotator/entityseg/mask2former/modeling/criterion.py

@@ -0,0 +1,263 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/models/detr.py
+"""
+MaskFormer criterion.
+"""
+import logging
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from detectron2.utils.comm import get_world_size
+from detectron2.projects.point_rend.point_features import (
+ get_uncertain_point_coords_with_randomness,
+ point_sample,
+)
+
+from ..utils.misc import is_dist_avail_and_initialized, nested_tensor_from_tensor_list
+
+
+def dice_loss(
+ inputs: torch.Tensor,
+ targets: torch.Tensor,
+ num_masks: float,
+ ):
+ """
+ Compute the DICE loss, similar to generalized IOU for masks
+ Args:
+ inputs: A float tensor of arbitrary shape.
+ The predictions for each example.
+ targets: A float tensor with the same shape as inputs. Stores the binary
+ classification label for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ """
+ inputs = inputs.sigmoid()
+ inputs = inputs.flatten(1)
+ numerator = 2 * (inputs * targets).sum(-1)
+ denominator = inputs.sum(-1) + targets.sum(-1)
+ loss = 1 - (numerator + 1) / (denominator + 1)
+ return loss.sum() / num_masks
+
+
+dice_loss_jit = torch.jit.script(
+ dice_loss
+) # type: torch.jit.ScriptModule
+
+
+def sigmoid_ce_loss(
+ inputs: torch.Tensor,
+ targets: torch.Tensor,
+ num_masks: float,
+ ):
+ """
+ Args:
+ inputs: A float tensor of arbitrary shape.
+ The predictions for each example.
+ targets: A float tensor with the same shape as inputs. Stores the binary
+ classification label for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ Returns:
+ Loss tensor
+ """
+ loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
+
+ return loss.mean(1).sum() / num_masks
+
+
+sigmoid_ce_loss_jit = torch.jit.script(
+ sigmoid_ce_loss
+) # type: torch.jit.ScriptModule
+
+
+def calculate_uncertainty(logits):
+ """
+ We estimate uncerainty as L1 distance between 0.0 and the logit prediction in 'logits' for the
+ foreground class in `classes`.
+ Args:
+ logits (Tensor): A tensor of shape (R, 1, ...) for class-specific or
+ class-agnostic, where R is the total number of predicted masks in all images and C is
+ the number of foreground classes. The values are logits.
+ Returns:
+ scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with
+ the most uncertain locations having the highest uncertainty score.
+ """
+ assert logits.shape[1] == 1
+ gt_class_logits = logits.clone()
+ return -(torch.abs(gt_class_logits))
+
+
+class SetCriterion(nn.Module):
+ """This class computes the loss for DETR.
+ The process happens in two steps:
+ 1) we compute hungarian assignment between ground truth boxes and the outputs of the model
+ 2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
+ """
+
+ def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses,
+ num_points, oversample_ratio, importance_sample_ratio):
+ """Create the criterion.
+ Parameters:
+ num_classes: number of object categories, omitting the special no-object category
+ matcher: module able to compute a matching between targets and proposals
+ weight_dict: dict containing as key the names of the losses and as values their relative weight.
+ eos_coef: relative classification weight applied to the no-object category
+ losses: list of all the losses to be applied. See get_loss for list of available losses.
+ """
+ super().__init__()
+ self.num_classes = num_classes
+ self.matcher = matcher
+ self.weight_dict = weight_dict
+ self.eos_coef = eos_coef
+ self.losses = losses
+ empty_weight = torch.ones(self.num_classes + 1)
+ empty_weight[-1] = self.eos_coef
+ self.register_buffer("empty_weight", empty_weight)
+
+ # pointwise mask loss parameters
+ self.num_points = num_points
+ self.oversample_ratio = oversample_ratio
+ self.importance_sample_ratio = importance_sample_ratio
+
+ def loss_labels(self, outputs, targets, indices, num_masks):
+ """Classification loss (NLL)
+ targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
+ """
+ assert "pred_logits" in outputs
+ src_logits = outputs["pred_logits"].float()
+
+ idx = self._get_src_permutation_idx(indices)
+ target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
+ target_classes = torch.full(
+ src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device
+ )
+ target_classes[idx] = target_classes_o
+
+ loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
+ losses = {"loss_ce": loss_ce}
+ return losses
+ 
+ def loss_masks(self, outputs, targets, indices, num_masks):
+ """Compute the losses related to the masks: the focal loss and the dice loss.
+ targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
+ """
+ assert "pred_masks" in outputs
+
+ src_idx = self._get_src_permutation_idx(indices)
+ tgt_idx = self._get_tgt_permutation_idx(indices)
+ src_masks = outputs["pred_masks"]
+ src_masks = src_masks[src_idx]
+ masks = [t["masks"] for t in targets]
+ # TODO use valid to mask invalid areas due to padding in loss
+ target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
+ target_masks = target_masks.to(src_masks)
+ target_masks = target_masks[tgt_idx]
+
+ # No need to upsample predictions as we are using normalized coordinates :)
+ # N x 1 x H x W
+ src_masks = src_masks[:, None]
+ target_masks = target_masks[:, None]
+
+ with torch.no_grad():
+ # sample point_coords
+ point_coords = get_uncertain_point_coords_with_randomness(
+ src_masks,
+ lambda logits: calculate_uncertainty(logits),
+ self.num_points,
+ self.oversample_ratio,
+ self.importance_sample_ratio,
+ )
+ # get gt labels
+ point_labels = point_sample(
+ target_masks,
+ point_coords,
+ align_corners=False,
+ ).squeeze(1)
+
+ point_logits = point_sample(
+ src_masks,
+ point_coords,
+ align_corners=False,
+ ).squeeze(1)
+
+ losses = {
+ "loss_mask": sigmoid_ce_loss_jit(point_logits, point_labels, num_masks),
+ "loss_dice": dice_loss_jit(point_logits, point_labels, num_masks),
+ }
+
+ del src_masks
+ del target_masks
+ return losses
+
+ def _get_src_permutation_idx(self, indices):
+ # permute predictions following indices
+ batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
+ src_idx = torch.cat([src for (src, _) in indices])
+ return batch_idx, src_idx
+
+ def _get_tgt_permutation_idx(self, indices):
+ # permute targets following indices
+ batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
+ tgt_idx = torch.cat([tgt for (_, tgt) in indices])
+ return batch_idx, tgt_idx
+
+ def get_loss(self, loss, outputs, targets, indices, num_masks):
+ loss_map = {
+ 'labels': self.loss_labels,
+ 'masks': self.loss_masks,
+ }
+ assert loss in loss_map, f"do you really want to compute {loss} loss?"
+ return loss_map[loss](outputs, targets, indices, num_masks)
+
+ def forward(self, outputs, targets):
+ """This performs the loss computation.
+ Parameters:
+ outputs: dict of tensors, see the output specification of the model for the format
+ targets: list of dicts, such that len(targets) == batch_size.
+ The expected keys in each dict depends on the losses applied, see each loss' doc
+ """
+ outputs_without_aux = {k: v for k, v in outputs.items() if k != "aux_outputs"}
+
+ # Retrieve the matching between the outputs of the last layer and the targets
+ indices = self.matcher(outputs_without_aux, targets)
+
+ # Compute the average number of target boxes accross all nodes, for normalization purposes
+ num_masks = sum(len(t["labels"]) for t in targets)
+ num_masks = torch.as_tensor(
+ [num_masks], dtype=torch.float, device=next(iter(outputs.values())).device
+ )
+ if is_dist_avail_and_initialized():
+ torch.distributed.all_reduce(num_masks)
+ num_masks = torch.clamp(num_masks / get_world_size(), min=1).item()
+
+ # Compute all the requested losses
+ losses = {}
+ for loss in self.losses:
+ losses.update(self.get_loss(loss, outputs, targets, indices, num_masks))
+
+ # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
+ if "aux_outputs" in outputs:
+ for i, aux_outputs in enumerate(outputs["aux_outputs"]):
+ indices = self.matcher(aux_outputs, targets)
+ for loss in self.losses:
+ l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_masks)
+ l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
+ losses.update(l_dict)
+
+ return losses
+
+ def __repr__(self):
+ head = "Criterion " + self.__class__.__name__
+ body = [
+ "matcher: {}".format(self.matcher.__repr__(_repr_indent=8)),
+ "losses: {}".format(self.losses),
+ "weight_dict: {}".format(self.weight_dict),
+ "num_classes: {}".format(self.num_classes),
+ "eos_coef: {}".format(self.eos_coef),
+ "num_points: {}".format(self.num_points),
+ "oversample_ratio: {}".format(self.oversample_ratio),
+ "importance_sample_ratio: {}".format(self.importance_sample_ratio),
+ ]
+ _repr_indent = 4
+ lines = [head] + [" " * _repr_indent + line for line in body]
+ return "\n".join(lines)

annotator/entityseg/mask2former/modeling/criterion_view.py

@@ -0,0 +1,288 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/models/detr.py
+"""
+MaskFormer criterion.
+"""
+import logging
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from detectron2.utils.comm import get_world_size
+from detectron2.projects.point_rend.point_features import (
+ get_uncertain_point_coords_with_randomness,
+ point_sample,
+)
+
+from mask2former.utils.misc import is_dist_avail_and_initialized
+
+import pdb
+
+
+def dice_loss(
+ inputs: torch.Tensor,
+ targets: torch.Tensor,
+ num_masks: float,
+ ):
+ """
+ Compute the DICE loss, similar to generalized IOU for masks
+ Args:
+ inputs: A float tensor of arbitrary shape.
+ The predictions for each example.
+ targets: A float tensor with the same shape as inputs. Stores the binary
+ classification label for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ """
+ inputs = inputs.sigmoid()
+ inputs = inputs.flatten(1)
+ numerator = 2 * (inputs * targets).sum(-1)
+ denominator = inputs.sum(-1) + targets.sum(-1)
+ loss = 1 - (numerator + 1) / (denominator + 1)
+ return loss.sum() / num_masks
+
+
+dice_loss_jit = torch.jit.script(
+ dice_loss
+) # type: torch.jit.ScriptModule
+
+
+def sigmoid_ce_loss(
+ inputs: torch.Tensor,
+ targets: torch.Tensor,
+ num_masks: float,
+ ):
+ """
+ Args:
+ inputs: A float tensor of arbitrary shape.
+ The predictions for each example.
+ targets: A float tensor with the same shape as inputs. Stores the binary
+ classification label for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ Returns:
+ Loss tensor
+ """
+ loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
+
+ return loss.mean(1).sum() / num_masks
+
+
+sigmoid_ce_loss_jit = torch.jit.script(
+ sigmoid_ce_loss
+) # type: torch.jit.ScriptModule
+
+
+def calculate_uncertainty(logits):
+ """
+ We estimate uncerainty as L1 distance between 0.0 and the logit prediction in 'logits' for the
+ foreground class in `classes`.
+ Args:
+ logits (Tensor): A tensor of shape (R, 1, ...) for class-specific or
+ class-agnostic, where R is the total number of predicted masks in all images and C is
+ the number of foreground classes. The values are logits.
+ Returns:
+ scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with
+ the most uncertain locations having the highest uncertainty score.
+ """
+ assert logits.shape[1] == 1
+ gt_class_logits = logits.clone()
+ return -(torch.abs(gt_class_logits))
+
+
+class ViewSetCriterion(nn.Module):
+ """This class computes the loss for DETR.
+ The process happens in two steps:
+ 1) we compute hungarian assignment between ground truth boxes and the outputs of the model
+ 2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
+ """
+
+ def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses,
+ num_points, oversample_ratio, importance_sample_ratio):
+ """Create the criterion.
+ Parameters:
+ num_classes: number of object categories, omitting the special no-object category
+ matcher: module able to compute a matching between targets and proposals
+ weight_dict: dict containing as key the names of the losses and as values their relative weight.
+ eos_coef: relative classification weight applied to the no-object category
+ losses: list of all the losses to be applied. See get_loss for list of available losses.
+ """
+ super().__init__()
+ self.num_classes = num_classes
+ self.matcher = matcher
+ self.weight_dict = weight_dict
+ self.eos_coef = eos_coef
+ self.losses = losses
+ empty_weight = torch.ones(self.num_classes + 1)
+ empty_weight[-1] = self.eos_coef
+ self.register_buffer("empty_weight", empty_weight)
+
+ # pointwise mask loss parameters
+ self.num_points = num_points
+ self.oversample_ratio = oversample_ratio
+ self.importance_sample_ratio = importance_sample_ratio
+
+ def loss_labels(self, outputs, targets, indices, num_masks):
+ """Classification loss (NLL)
+ targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
+ """
+ assert "pred_logits" in outputs
+ src_logits = outputs["pred_logits"].float()
+ ## src_logits: torch.Size([2, 100, 41])
+
+ idx = self._get_src_permutation_idx(indices)
+ ## idx: (tensor([0, 0, 1, 1]), tensor([17, 84, 17, 76]))
+ target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
+ ### target_class_o: tensor([ 0, 26, 0, 11], device='cuda:0')
+ target_classes = torch.full(
+ src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device
+ )
+ ## target_class: torch.Size([2, 100]), 全是40, background类
+ target_classes[idx] = target_classes_o
+ ## 
+ ## src_logits: torch.Size([2, 41, 100])
+ ## target_classes: torch.Size([2, 100])
+ ## self.empty_weight: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
+ ##1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
+ ##1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
+ ##1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
+ ##1.0000, 1.0000, 1.0000, 1.0000, 0.1000], device='cuda:0')
+ loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
+ losses = {"loss_ce": loss_ce}
+ return losses
+ 
+ def loss_masks(self, outputs, targets, indices, num_masks):
+ """Compute the losses related to the masks: the focal loss and the dice loss.
+ targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
+ """
+ assert "pred_masks" in outputs
+
+ src_idx = self._get_src_permutation_idx(indices)
+ ### src_idx: (tensor([0, 0, 1, 1]), tensor([34, 95, 32, 65]))
+ src_masks = outputs["pred_masks"]
+ ## src_masks: torch.Size([2, 100, 2, 120, 216])
+ src_masks = src_masks[src_idx]
+ ## src_masks: torch.Size([4, 2, 120, 216])
+
+ # Modified to handle video
+ target_masks = torch.cat([t['masks'][i] for t, (_, i) in zip(targets, indices)]).to(src_masks)
+ ### target_masks: torch.Size([4, 2, 480, 864])
+
+ # No need to upsample predictions as we are using normalized coordinates :)
+ # NT x 1 x H x W
+ src_masks = src_masks.flatten(0, 1)[:, None]
+ ## src_masks: torch.Size([8, 1, 120, 216])
+ target_masks = target_masks.flatten(0, 1)[:, None]
+ ## target_masks: torch.Size([8, 1, 480, 864])
+
+ with torch.no_grad():
+ # sample point_coords
+ point_coords = get_uncertain_point_coords_with_randomness(
+ src_masks,
+ lambda logits: calculate_uncertainty(logits),
+ self.num_points,
+ self.oversample_ratio,
+ self.importance_sample_ratio,
+ )
+ # get gt labels
+ point_labels = point_sample(
+ target_masks,
+ point_coords,
+ align_corners=False,
+ ).squeeze(1)
+
+ point_logits = point_sample(
+ src_masks,
+ point_coords,
+ align_corners=False,
+ ).squeeze(1)
+
+ losses = {
+ "loss_mask": sigmoid_ce_loss_jit(point_logits, point_labels, num_masks),
+ "loss_dice": dice_loss_jit(point_logits, point_labels, num_masks),
+ }
+
+ del src_masks
+ del target_masks
+ return losses
+
+ def _get_src_permutation_idx(self, indices):
+ # permute predictions following indices
+ batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
+ src_idx = torch.cat([src for (src, _) in indices])
+ return batch_idx, src_idx
+
+ def _get_tgt_permutation_idx(self, indices):
+ # permute targets following indices
+ batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
+ tgt_idx = torch.cat([tgt for (_, tgt) in indices])
+ return batch_idx, tgt_idx
+
+ def get_loss(self, loss, outputs, targets, indices, num_masks):
+ loss_map = {
+ 'labels': self.loss_labels,
+ 'masks': self.loss_masks,
+ }
+ assert loss in loss_map, f"do you really want to compute {loss} loss?"
+ return loss_map[loss](outputs, targets, indices, num_masks)
+
+ def forward(self, outputs, targets, return_indices=False):
+ """This performs the loss computation.
+ Parameters:
+ outputs: dict of tensors, see the output specification of the model for the format
+ targets: list of dicts, such that len(targets) == batch_size.
+ The expected keys in each dict depends on the losses applied, see each loss' doc
+ """
+ outputs_without_aux = {k: v for k, v in outputs.items() if k != "aux_outputs"}
+
+ # Retrieve the matching between the outputs of the last layer and the targets
+ indices = self.matcher(outputs_without_aux, targets)
+ indices_l = []
+ indices_l.append(indices)
+ # pdb.set_trace()
+
+ # Compute the average number of target boxes accross all nodes, for normalization purposes
+ num_masks = sum(len(t["labels"]) for t in targets)
+ num_masks = torch.as_tensor(
+ [num_masks], dtype=torch.float, device=next(iter(outputs.values())).device
+ )
+ if is_dist_avail_and_initialized():
+ torch.distributed.all_reduce(num_masks)
+ num_masks = torch.clamp(num_masks / get_world_size(), min=1).item()
+
+ # Compute all the requested losses
+ losses = {}
+ for loss in self.losses:
+ losses.update(self.get_loss(loss, outputs, targets, indices, num_masks))
+
+ # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
+ if "aux_outputs" in outputs:
+ for i, aux_outputs in enumerate(outputs["aux_outputs"]):
+ indices = self.matcher(aux_outputs, targets)
+ indices_l.append(indices)
+ for loss in self.losses:
+ l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_masks)
+ l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
+ losses.update(l_dict)
+ indices_l.append(indices_l[0])
+ indices_l = indices_l[1:]
+
+ if return_indices:
+ return losses, indices_l
+ else:
+ return losses
+
+ def __repr__(self):
+ head = "Criterion " + self.__class__.__name__
+ body = [
+ "matcher: {}".format(self.matcher.__repr__(_repr_indent=8)),
+ "losses: {}".format(self.losses),
+ "weight_dict: {}".format(self.weight_dict),
+ "num_classes: {}".format(self.num_classes),
+ "eos_coef: {}".format(self.eos_coef),
+ "num_points: {}".format(self.num_points),
+ "oversample_ratio: {}".format(self.oversample_ratio),
+ "importance_sample_ratio: {}".format(self.importance_sample_ratio),
+ ]
+ _repr_indent = 4
+ lines = [head] + [" " * _repr_indent + line for line in body]
+ return "\n".join(lines)

annotator/entityseg/mask2former/modeling/matcher.py

@@ -0,0 +1,189 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/models/matcher.py
+"""
+Modules to compute the matching cost and solve the corresponding LSAP.
+"""
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+from torch import nn
+from torch.cuda.amp import autocast
+
+from detectron2.projects.point_rend.point_features import point_sample
+
+
+def batch_dice_loss(inputs: torch.Tensor, targets: torch.Tensor):
+ """
+ Compute the DICE loss, similar to generalized IOU for masks
+ Args:
+ inputs: A float tensor of arbitrary shape.
+ The predictions for each example.
+ targets: A float tensor with the same shape as inputs. Stores the binary
+ classification label for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ """
+ inputs = inputs.sigmoid()
+ inputs = inputs.flatten(1)
+ numerator = 2 * torch.einsum("nc,mc->nm", inputs, targets)
+ denominator = inputs.sum(-1)[:, None] + targets.sum(-1)[None, :]
+ loss = 1 - (numerator + 1) / (denominator + 1)
+ return loss
+
+
+batch_dice_loss_jit = torch.jit.script(
+ batch_dice_loss
+) # type: torch.jit.ScriptModule
+
+
+def batch_sigmoid_ce_loss(inputs: torch.Tensor, targets: torch.Tensor):
+ """
+ Args:
+ inputs: A float tensor of arbitrary shape.
+ The predictions for each example.
+ targets: A float tensor with the same shape as inputs. Stores the binary
+ classification label for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ Returns:
+ Loss tensor
+ """
+ hw = inputs.shape[1]
+
+ pos = F.binary_cross_entropy_with_logits(
+ inputs, torch.ones_like(inputs), reduction="none"
+ )
+ neg = F.binary_cross_entropy_with_logits(
+ inputs, torch.zeros_like(inputs), reduction="none"
+ )
+
+ loss = torch.einsum("nc,mc->nm", pos, targets) + torch.einsum(
+ "nc,mc->nm", neg, (1 - targets)
+ )
+
+ return loss / hw
+
+
+batch_sigmoid_ce_loss_jit = torch.jit.script(
+ batch_sigmoid_ce_loss
+) # type: torch.jit.ScriptModule
+
+
+class HungarianMatcher(nn.Module):
+ """This class computes an assignment between the targets and the predictions of the network
+
+ For efficiency reasons, the targets don't include the no_object. Because of this, in general,
+ there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
+ while the others are un-matched (and thus treated as non-objects).
+ """
+
+ def __init__(self, cost_class: float = 1, cost_mask: float = 1, cost_dice: float = 1, num_points: int = 0):
+ """Creates the matcher
+
+ Params:
+ cost_class: This is the relative weight of the classification error in the matching cost
+ cost_mask: This is the relative weight of the focal loss of the binary mask in the matching cost
+ cost_dice: This is the relative weight of the dice loss of the binary mask in the matching cost
+ """
+ super().__init__()
+ self.cost_class = cost_class
+ self.cost_mask = cost_mask
+ self.cost_dice = cost_dice
+
+ assert cost_class != 0 or cost_mask != 0 or cost_dice != 0, "all costs cant be 0"
+
+ self.num_points = num_points
+
+ @torch.no_grad()
+ def memory_efficient_forward(self, outputs, targets):
+ """More memory-friendly matching"""
+ bs, num_queries = outputs["pred_logits"].shape[:2]
+
+ indices = []
+
+ # Iterate through batch size
+ for b in range(bs):
+
+ out_prob = outputs["pred_logits"][b].softmax(-1) # [num_queries, num_classes]
+ tgt_ids = targets[b]["labels"]
+
+ # Compute the classification cost. Contrary to the loss, we don't use the NLL,
+ # but approximate it in 1 - proba[target class].
+ # The 1 is a constant that doesn't change the matching, it can be ommitted.
+ cost_class = -out_prob[:, tgt_ids]
+
+ out_mask = outputs["pred_masks"][b] # [num_queries, H_pred, W_pred]
+ # gt masks are already padded when preparing target
+ tgt_mask = targets[b]["masks"].to(out_mask)
+
+ out_mask = out_mask[:, None]
+ tgt_mask = tgt_mask[:, None]
+ # all masks share the same set of points for efficient matching!
+ point_coords = torch.rand(1, self.num_points, 2, device=out_mask.device)
+ # get gt labels
+ tgt_mask = point_sample(
+ tgt_mask,
+ point_coords.repeat(tgt_mask.shape[0], 1, 1),
+ align_corners=False,
+ ).squeeze(1)
+
+ out_mask = point_sample(
+ out_mask,
+ point_coords.repeat(out_mask.shape[0], 1, 1),
+ align_corners=False,
+ ).squeeze(1)
+
+ with autocast(enabled=False):
+ out_mask = out_mask.float()
+ tgt_mask = tgt_mask.float()
+ # Compute the focal loss between masks
+ cost_mask = batch_sigmoid_ce_loss(out_mask, tgt_mask)
+
+ # Compute the dice loss betwen masks
+ cost_dice = batch_dice_loss(out_mask, tgt_mask)
+ 
+ # Final cost matrix
+ C = (
+ self.cost_mask * cost_mask
+ + self.cost_class * cost_class
+ + self.cost_dice * cost_dice
+ )
+ C = C.reshape(num_queries, -1).cpu()
+
+ indices.append(linear_sum_assignment(C))
+
+ return [
+ (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64))
+ for i, j in indices
+ ]
+
+ @torch.no_grad()
+ def forward(self, outputs, targets):
+ """Performs the matching
+
+ Params:
+ outputs: This is a dict that contains at least these entries:
+ "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
+ "pred_masks": Tensor of dim [batch_size, num_queries, H_pred, W_pred] with the predicted masks
+
+ targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
+ "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth
+ objects in the target) containing the class labels
+ "masks": Tensor of dim [num_target_boxes, H_gt, W_gt] containing the target masks
+
+ Returns:
+ A list of size batch_size, containing tuples of (index_i, index_j) where:
+ - index_i is the indices of the selected predictions (in order)
+ - index_j is the indices of the corresponding selected targets (in order)
+ For each batch element, it holds:
+ len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
+ """
+ return self.memory_efficient_forward(outputs, targets)
+
+ def __repr__(self, _repr_indent=4):
+ head = "Matcher " + self.__class__.__name__
+ body = [
+ "cost_class: {}".format(self.cost_class),
+ "cost_mask: {}".format(self.cost_mask),
+ "cost_dice: {}".format(self.cost_dice),
+ ]
+ lines = [head] + [" " * _repr_indent + line for line in body]
+ return "\n".join(lines)

annotator/entityseg/mask2former/modeling/matcher_view.py

@@ -0,0 +1,194 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/models/matcher.py
+"""
+Modules to compute the matching cost and solve the corresponding LSAP.
+"""
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+from torch import nn
+from torch.cuda.amp import autocast
+
+from detectron2.projects.point_rend.point_features import point_sample
+
+def batch_dice_loss(inputs: torch.Tensor, targets: torch.Tensor):
+ """
+ Compute the DICE loss, similar to generalized IOU for masks
+ Args:
+ inputs: A float tensor of arbitrary shape.
+ The predictions for each example.
+ targets: A float tensor with the same shape as inputs. Stores the binary
+ classification label for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ """
+ inputs = inputs.sigmoid()
+ inputs = inputs.flatten(1)
+ numerator = 2 * torch.einsum("nc,mc->nm", inputs, targets)
+ denominator = inputs.sum(-1)[:, None] + targets.sum(-1)[None, :]
+ loss = 1 - (numerator + 1) / (denominator + 1)
+ return loss
+
+
+batch_dice_loss_jit = torch.jit.script(
+ batch_dice_loss
+) # type: torch.jit.ScriptModule
+
+
+def batch_sigmoid_ce_loss(inputs: torch.Tensor, targets: torch.Tensor):
+ """
+ Args:
+ inputs: A float tensor of arbitrary shape.
+ The predictions for each example.
+ targets: A float tensor with the same shape as inputs. Stores the binary
+ classification label for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ Returns:
+ Loss tensor
+ """
+ hw = inputs.shape[1]
+
+ pos = F.binary_cross_entropy_with_logits(
+ inputs, torch.ones_like(inputs), reduction="none"
+ )
+ neg = F.binary_cross_entropy_with_logits(
+ inputs, torch.zeros_like(inputs), reduction="none"
+ )
+
+ loss = torch.einsum("nc,mc->nm", pos, targets) + torch.einsum(
+ "nc,mc->nm", neg, (1 - targets)
+ )
+
+ return loss / hw
+
+
+batch_sigmoid_ce_loss_jit = torch.jit.script(
+ batch_sigmoid_ce_loss
+) # type: torch.jit.ScriptModule
+
+
+class ViewHungarianMatcher(nn.Module):
+ """This class computes an assignment between the targets and the predictions of the network
+
+ For efficiency reasons, the targets don't include the no_object. Because of this, in general,
+ there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
+ while the others are un-matched (and thus treated as non-objects).
+ """
+
+ def __init__(self, cost_class: float = 1, cost_mask: float = 1, cost_dice: float = 1, num_points: int = 0):
+ """Creates the matcher
+
+ Params:
+ cost_class: This is the relative weight of the classification error in the matching cost
+ cost_mask: This is the relative weight of the focal loss of the binary mask in the matching cost
+ cost_dice: This is the relative weight of the dice loss of the binary mask in the matching cost
+ """
+ super().__init__()
+ self.cost_class = cost_class
+ self.cost_mask = cost_mask
+ self.cost_dice = cost_dice
+
+ assert cost_class != 0 or cost_mask != 0 or cost_dice != 0, "all costs cant be 0"
+
+ self.num_points = num_points
+
+ @torch.no_grad()
+ def memory_efficient_forward(self, outputs, targets):
+ """More memory-friendly matching"""
+ ### outputs["pred_logits"]: torch.Size([2, 100, 41]), query是对两帧负责，所以没有frame的概念
+ ### outputs["pred_masks"]: torch.Size([2, 100, 2, 120, 160]), 第三维的2是两帧frame
+ bs, num_queries = outputs["pred_logits"].shape[:2]
+
+ indices = []
+
+ # Iterate through batch size
+ for b in range(bs):
+ out_prob = outputs["pred_logits"][b].softmax(-1) # [num_queries, num_classes]
+ ## out_prob: [100, 41], 100个query, 40类+background类
+ tgt_ids = targets[b]["labels"]
+ ## tgt_ids: tensor([ 3, 10]), 说明只有两个ground truth
+
+ # Compute the classification cost. Contrary to the loss, we don't use the NLL,
+ # but approximate it in 1 - proba[target class].
+ # The 1 is a constant that doesn't change the matching, it can be ommitted.
+ cost_class = -out_prob[:, tgt_ids]
+
+ out_mask = outputs["pred_masks"][b] # [num_queries, T, H_pred, W_pred]
+ ### out_mask: torch.Size([100, 2, 120, 160])
+ # gt masks are already padded when preparing target
+ tgt_mask = targets[b]["masks"].to(out_mask) # [num_gts, T, H_pred, W_pred]
+ ## tgt_mask: torch.Size([2, 2, 480, 640])
+
+ # out_mask = out_mask[:, None]
+ # tgt_mask = tgt_mask[:, None]
+ # all masks share the same set of points for efficient matching!
+ point_coords = torch.rand(1, self.num_points, 2, device=out_mask.device)
+ # get gt labels
+ tgt_mask = point_sample(
+ tgt_mask,
+ point_coords.repeat(tgt_mask.shape[0], 1, 1), ## repeat了一份, torch.Size([2, 12544, 2]), 每一帧采样的位置都是一样的
+ align_corners=False,
+ ).flatten(1)
+
+ out_mask = point_sample(
+ out_mask,
+ point_coords.repeat(out_mask.shape[0], 1, 1),
+ align_corners=False,
+ ).flatten(1)
+
+ with autocast(enabled=False):
+ out_mask = out_mask.float() ## out_mask: torch.Size([100, 25088])
+ tgt_mask = tgt_mask.float() ## tgt_mask: torch.Size([2, 25088])
+ # Compute the focal loss between masks
+ cost_mask = batch_sigmoid_ce_loss_jit(out_mask, tgt_mask) ## cost_mask: torch.Size([100, 2])
+
+ # Compute the dice loss betwen masks
+ cost_dice = batch_dice_loss_jit(out_mask, tgt_mask) ## cost_dice: torch.Size([100, 2])
+ 
+ # Final cost matrix
+ C = (
+ self.cost_mask * cost_mask
+ + self.cost_class * cost_class
+ + self.cost_dice * cost_dice
+ )
+ C = C.reshape(num_queries, -1).cpu()
+
+ indices.append(linear_sum_assignment(C))
+ ## [(array([17, 33]), array([1, 0]), ...]
+
+ return [
+ (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64))
+ for i, j in indices
+ ]
+
+ @torch.no_grad()
+ def forward(self, outputs, targets):
+ """Performs the matching
+
+ Params:
+ outputs: This is a dict that contains at least these entries:
+ "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
+ "pred_masks": Tensor of dim [batch_size, num_queries, H_pred, W_pred] with the predicted masks
+
+ targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
+ "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth
+ objects in the target) containing the class labels
+ "masks": Tensor of dim [num_target_boxes, H_gt, W_gt] containing the target masks
+
+ Returns:
+ A list of size batch_size, containing tuples of (index_i, index_j) where:
+ - index_i is the indices of the selected predictions (in order)
+ - index_j is the indices of the corresponding selected targets (in order)
+ For each batch element, it holds:
+ len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
+ """
+ return self.memory_efficient_forward(outputs, targets)
+
+ def __repr__(self, _repr_indent=4):
+ head = "Matcher " + self.__class__.__name__
+ body = [
+ "cost_class: {}".format(self.cost_class),
+ "cost_mask: {}".format(self.cost_mask),
+ "cost_dice: {}".format(self.cost_dice),
+ ]
+ lines = [head] + [" " * _repr_indent + line for line in body]
+ return "\n".join(lines)

annotator/entityseg/mask2former/modeling/meta_arch/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Facebook, Inc. and its affiliates.

annotator/entityseg/mask2former/modeling/meta_arch/mask_former_head.py

@@ -0,0 +1,133 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import logging
+from copy import deepcopy
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import fvcore.nn.weight_init as weight_init
+from torch import nn
+from torch.nn import functional as F
+
+from detectron2.config import configurable
+from detectron2.layers import Conv2d, ShapeSpec, get_norm
+from detectron2.modeling import SEM_SEG_HEADS_REGISTRY
+
+from ..transformer_decoder.maskformer_transformer_decoder import build_transformer_decoder
+from ..pixel_decoder.fpn import build_pixel_decoder
+
+
+@SEM_SEG_HEADS_REGISTRY.register()
+class MaskFormerHead(nn.Module):
+
+ _version = 2
+
+ def _load_from_state_dict(
+ self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+ ):
+ version = local_metadata.get("version", None)
+ if version is None or version < 2:
+ # Do not warn if train from scratch
+ scratch = True
+ logger = logging.getLogger(__name__)
+ for k in list(state_dict.keys()):
+ newk = k
+ if "sem_seg_head" in k and not k.startswith(prefix + "predictor"):
+ # newk = k.replace(prefix, prefix + "pixel_decoder.")
+ newk = k.replace(prefix, prefix)
+ # logger.debug(f"{k} ==> {newk}")
+ if newk != k:
+ state_dict[newk] = state_dict[k]
+ del state_dict[k]
+ scratch = False
+
+ if not scratch:
+ logger.warning(
+ f"Weight format of {self.__class__.__name__} have changed! "
+ "Please upgrade your models. Applying automatic conversion now ..."
+ )
+
+ @configurable
+ def __init__(
+ self,
+ input_shape: Dict[str, ShapeSpec],
+ *,
+ num_classes: int,
+ pixel_decoder: nn.Module,
+ loss_weight: float = 1.0,
+ ignore_value: int = -1,
+ # extra parameters
+ transformer_predictor: nn.Module,
+ transformer_in_feature: str,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ input_shape: shapes (channels and stride) of the input features
+ num_classes: number of classes to predict
+ pixel_decoder: the pixel decoder module
+ loss_weight: loss weight
+ ignore_value: category id to be ignored during training.
+ transformer_predictor: the transformer decoder that makes prediction
+ transformer_in_feature: input feature name to the transformer_predictor
+ """
+ super().__init__()
+ input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
+ self.in_features = [k for k, v in input_shape]
+ feature_strides = [v.stride for k, v in input_shape]
+ feature_channels = [v.channels for k, v in input_shape]
+
+ self.ignore_value = ignore_value
+ self.common_stride = 4
+ self.loss_weight = loss_weight
+
+ self.pixel_decoder = pixel_decoder
+ self.predictor = transformer_predictor
+ self.transformer_in_feature = transformer_in_feature
+
+ self.num_classes = num_classes
+
+ @classmethod
+ def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
+ # figure out in_channels to transformer predictor
+ if cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "transformer_encoder":
+ transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
+ elif cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "pixel_embedding":
+ transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.MASK_DIM
+ elif cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "multi_scale_pixel_decoder": # for maskformer2
+ transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
+ else:
+ transformer_predictor_in_channels = input_shape[cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE].channels
+
+ return {
+ "input_shape": {
+ k: v for k, v in input_shape.items() if k in cfg.MODEL.SEM_SEG_HEAD.IN_FEATURES
+ },
+ "ignore_value": cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
+ "num_classes": cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
+ "pixel_decoder": build_pixel_decoder(cfg, input_shape),
+ "loss_weight": cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT,
+ "transformer_in_feature": cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE,
+ "transformer_predictor": build_transformer_decoder(
+ cfg,
+ transformer_predictor_in_channels,
+ mask_classification=True,
+ ),
+ }
+
+ def forward(self, features, mask=None):
+ return self.layers(features, mask)
+
+ def layers(self, features, mask=None):
+ mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
+ if self.transformer_in_feature == "multi_scale_pixel_decoder":
+ predictions = self.predictor(multi_scale_features, mask_features, mask)
+ else:
+ if self.transformer_in_feature == "transformer_encoder":
+ assert (
+ transformer_encoder_features is not None
+ ), "Please use the TransformerEncoderPixelDecoder."
+ predictions = self.predictor(transformer_encoder_features, mask_features, mask)
+ elif self.transformer_in_feature == "pixel_embedding":
+ predictions = self.predictor(mask_features, mask_features, mask)
+ else:
+ predictions = self.predictor(features[self.transformer_in_feature], mask_features, mask)
+ return predictions

annotator/entityseg/mask2former/modeling/meta_arch/per_pixel_baseline.py

@@ -0,0 +1,243 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import logging
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import fvcore.nn.weight_init as weight_init
+from torch import nn
+from torch.nn import functional as F
+
+from detectron2.config import configurable
+from detectron2.layers import Conv2d, ShapeSpec, get_norm
+from detectron2.modeling import SEM_SEG_HEADS_REGISTRY
+
+from ..transformer_decoder.maskformer_transformer_decoder import StandardTransformerDecoder
+from ..pixel_decoder.fpn import build_pixel_decoder
+
+
+@SEM_SEG_HEADS_REGISTRY.register()
+class PerPixelBaselineHead(nn.Module):
+
+ _version = 2
+
+ def _load_from_state_dict(
+ self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+ ):
+ version = local_metadata.get("version", None)
+ if version is None or version < 2:
+ logger = logging.getLogger(__name__)
+ # Do not warn if train from scratch
+ scratch = True
+ logger = logging.getLogger(__name__)
+ for k in list(state_dict.keys()):
+ newk = k
+ if "sem_seg_head" in k and not k.startswith(prefix + "predictor"):
+ newk = k.replace(prefix, prefix + "pixel_decoder.")
+ # logger.warning(f"{k} ==> {newk}")
+ if newk != k:
+ state_dict[newk] = state_dict[k]
+ del state_dict[k]
+ scratch = False
+
+ if not scratch:
+ logger.warning(
+ f"Weight format of {self.__class__.__name__} have changed! "
+ "Please upgrade your models. Applying automatic conversion now ..."
+ )
+
+ @configurable
+ def __init__(
+ self,
+ input_shape: Dict[str, ShapeSpec],
+ *,
+ num_classes: int,
+ pixel_decoder: nn.Module,
+ loss_weight: float = 1.0,
+ ignore_value: int = -1,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ input_shape: shapes (channels and stride) of the input features
+ num_classes: number of classes to predict
+ pixel_decoder: the pixel decoder module
+ loss_weight: loss weight
+ ignore_value: category id to be ignored during training.
+ """
+ super().__init__()
+ input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
+ self.in_features = [k for k, v in input_shape]
+ feature_strides = [v.stride for k, v in input_shape]
+ feature_channels = [v.channels for k, v in input_shape]
+
+ self.ignore_value = ignore_value
+ self.common_stride = 4
+ self.loss_weight = loss_weight
+
+ self.pixel_decoder = pixel_decoder
+ self.predictor = Conv2d(
+ self.pixel_decoder.mask_dim, num_classes, kernel_size=1, stride=1, padding=0
+ )
+ weight_init.c2_msra_fill(self.predictor)
+
+ @classmethod
+ def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
+ return {
+ "input_shape": {
+ k: v for k, v in input_shape.items() if k in cfg.MODEL.SEM_SEG_HEAD.IN_FEATURES
+ },
+ "ignore_value": cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
+ "num_classes": cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
+ "pixel_decoder": build_pixel_decoder(cfg, input_shape),
+ "loss_weight": cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT,
+ }
+
+ def forward(self, features, targets=None):
+ """
+ Returns:
+ In training, returns (None, dict of losses)
+ In inference, returns (CxHxW logits, {})
+ """
+ x = self.layers(features)
+ if self.training:
+ return None, self.losses(x, targets)
+ else:
+ x = F.interpolate(
+ x, scale_factor=self.common_stride, mode="bilinear", align_corners=False
+ )
+ return x, {}
+
+ def layers(self, features):
+ x, _, _ = self.pixel_decoder.forward_features(features)
+ x = self.predictor(x)
+ return x
+
+ def losses(self, predictions, targets):
+ predictions = predictions.float() # https://github.com/pytorch/pytorch/issues/48163
+ predictions = F.interpolate(
+ predictions, scale_factor=self.common_stride, mode="bilinear", align_corners=False
+ )
+ loss = F.cross_entropy(
+ predictions, targets, reduction="mean", ignore_index=self.ignore_value
+ )
+ losses = {"loss_sem_seg": loss * self.loss_weight}
+ return losses
+
+
+@SEM_SEG_HEADS_REGISTRY.register()
+class PerPixelBaselinePlusHead(PerPixelBaselineHead):
+ def _load_from_state_dict(
+ self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+ ):
+ version = local_metadata.get("version", None)
+ if version is None or version < 2:
+ # Do not warn if train from scratch
+ scratch = True
+ logger = logging.getLogger(__name__)
+ for k in list(state_dict.keys()):
+ newk = k
+ if "sem_seg_head" in k and not k.startswith(prefix + "predictor"):
+ newk = k.replace(prefix, prefix + "pixel_decoder.")
+ logger.debug(f"{k} ==> {newk}")
+ if newk != k:
+ state_dict[newk] = state_dict[k]
+ del state_dict[k]
+ scratch = False
+
+ if not scratch:
+ logger.warning(
+ f"Weight format of {self.__class__.__name__} have changed! "
+ "Please upgrade your models. Applying automatic conversion now ..."
+ )
+
+ @configurable
+ def __init__(
+ self,
+ input_shape: Dict[str, ShapeSpec],
+ *,
+ # extra parameters
+ transformer_predictor: nn.Module,
+ transformer_in_feature: str,
+ deep_supervision: bool,
+ # inherit parameters
+ num_classes: int,
+ pixel_decoder: nn.Module,
+ loss_weight: float = 1.0,
+ ignore_value: int = -1,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ input_shape: shapes (channels and stride) of the input features
+ transformer_predictor: the transformer decoder that makes prediction
+ transformer_in_feature: input feature name to the transformer_predictor
+ deep_supervision: whether or not to add supervision to the output of
+ every transformer decoder layer
+ num_classes: number of classes to predict
+ pixel_decoder: the pixel decoder module
+ loss_weight: loss weight
+ ignore_value: category id to be ignored during training.
+ """
+ super().__init__(
+ input_shape,
+ num_classes=num_classes,
+ pixel_decoder=pixel_decoder,
+ loss_weight=loss_weight,
+ ignore_value=ignore_value,
+ )
+
+ del self.predictor
+
+ self.predictor = transformer_predictor
+ self.transformer_in_feature = transformer_in_feature
+ self.deep_supervision = deep_supervision
+
+ @classmethod
+ def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
+ ret = super().from_config(cfg, input_shape)
+ ret["transformer_in_feature"] = cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE
+ if cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "transformer_encoder":
+ in_channels = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
+ else:
+ in_channels = input_shape[ret["transformer_in_feature"]].channels
+ ret["transformer_predictor"] = StandardTransformerDecoder(
+ cfg, in_channels, mask_classification=False
+ )
+ ret["deep_supervision"] = cfg.MODEL.MASK_FORMER.DEEP_SUPERVISION
+ return ret
+
+ def forward(self, features, targets=None):
+ """
+ Returns:
+ In training, returns (None, dict of losses)
+ In inference, returns (CxHxW logits, {})
+ """
+ x, aux_outputs = self.layers(features)
+ if self.training:
+ if self.deep_supervision:
+ losses = self.losses(x, targets)
+ for i, aux_output in enumerate(aux_outputs):
+ losses["loss_sem_seg" + f"_{i}"] = self.losses(
+ aux_output["pred_masks"], targets
+ )["loss_sem_seg"]
+ return None, losses
+ else:
+ return None, self.losses(x, targets)
+ else:
+ x = F.interpolate(
+ x, scale_factor=self.common_stride, mode="bilinear", align_corners=False
+ )
+ return x, {}
+
+ def layers(self, features):
+ mask_features, transformer_encoder_features, _ = self.pixel_decoder.forward_features(features)
+ if self.transformer_in_feature == "transformer_encoder":
+ assert (
+ transformer_encoder_features is not None
+ ), "Please use the TransformerEncoderPixelDecoder."
+ predictions = self.predictor(transformer_encoder_features, mask_features)
+ else:
+ predictions = self.predictor(features[self.transformer_in_feature], mask_features)
+ if self.deep_supervision:
+ return predictions["pred_masks"], predictions["aux_outputs"]
+ else:
+ return predictions["pred_masks"], None

annotator/entityseg/mask2former/modeling/pixel_decoder/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Facebook, Inc. and its affiliates.

annotator/entityseg/mask2former/modeling/pixel_decoder/fpn.py

@@ -0,0 +1,312 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import logging
+import numpy as np
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import fvcore.nn.weight_init as weight_init
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.init import xavier_uniform_, constant_, uniform_, normal_
+from torch.cuda.amp import autocast
+
+from detectron2.config import configurable
+from detectron2.layers import Conv2d, DeformConv, ShapeSpec, get_norm
+from detectron2.modeling import SEM_SEG_HEADS_REGISTRY
+
+from ..transformer_decoder.position_encoding import PositionEmbeddingSine
+from ..transformer_decoder.transformer import TransformerEncoder, TransformerEncoderLayer, _get_clones, _get_activation_fn
+
+
+def build_pixel_decoder(cfg, input_shape):
+ """
+ Build a pixel decoder from `cfg.MODEL.MASK_FORMER.PIXEL_DECODER_NAME`.
+ """
+ name = cfg.MODEL.SEM_SEG_HEAD.PIXEL_DECODER_NAME
+ model = SEM_SEG_HEADS_REGISTRY.get(name)(cfg, input_shape)
+ forward_features = getattr(model, "forward_features", None)
+ if not callable(forward_features):
+ raise ValueError(
+ "Only SEM_SEG_HEADS with forward_features method can be used as pixel decoder. "
+ f"Please implement forward_features for {name} to only return mask features."
+ )
+ return model
+
+
+# This is a modified FPN decoder.
+@SEM_SEG_HEADS_REGISTRY.register()
+class BasePixelDecoder(nn.Module):
+ @configurable
+ def __init__(
+ self,
+ input_shape: Dict[str, ShapeSpec],
+ *,
+ conv_dim: int,
+ mask_dim: int,
+ norm: Optional[Union[str, Callable]] = None,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ input_shape: shapes (channels and stride) of the input features
+ conv_dims: number of output channels for the intermediate conv layers.
+ mask_dim: number of output channels for the final conv layer.
+ norm (str or callable): normalization for all conv layers
+ """
+ super().__init__()
+
+ input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
+ self.in_features = [k for k, v in input_shape] # starting from "res2" to "res5"
+ feature_channels = [v.channels for k, v in input_shape]
+
+ lateral_convs = []
+ output_convs = []
+
+ use_bias = norm == ""
+ for idx, in_channels in enumerate(feature_channels):
+ if idx == len(self.in_features) - 1:
+ output_norm = get_norm(norm, conv_dim)
+ output_conv = Conv2d(
+ in_channels,
+ conv_dim,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ bias=use_bias,
+ norm=output_norm,
+ activation=F.relu,
+ )
+ weight_init.c2_xavier_fill(output_conv)
+ self.add_module("layer_{}".format(idx + 1), output_conv)
+
+ lateral_convs.append(None)
+ output_convs.append(output_conv)
+ else:
+ lateral_norm = get_norm(norm, conv_dim)
+ output_norm = get_norm(norm, conv_dim)
+
+ lateral_conv = Conv2d(
+ in_channels, conv_dim, kernel_size=1, bias=use_bias, norm=lateral_norm
+ )
+ output_conv = Conv2d(
+ conv_dim,
+ conv_dim,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ bias=use_bias,
+ norm=output_norm,
+ activation=F.relu,
+ )
+ weight_init.c2_xavier_fill(lateral_conv)
+ weight_init.c2_xavier_fill(output_conv)
+ self.add_module("adapter_{}".format(idx + 1), lateral_conv)
+ self.add_module("layer_{}".format(idx + 1), output_conv)
+
+ lateral_convs.append(lateral_conv)
+ output_convs.append(output_conv)
+ # Place convs into top-down order (from low to high resolution)
+ # to make the top-down computation in forward clearer.
+ self.lateral_convs = lateral_convs[::-1]
+ self.output_convs = output_convs[::-1]
+
+ self.mask_dim = mask_dim
+ self.mask_features = Conv2d(
+ conv_dim,
+ mask_dim,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+ weight_init.c2_xavier_fill(self.mask_features)
+
+ self.maskformer_num_feature_levels = 3 # always use 3 scales
+
+ @classmethod
+ def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
+ ret = {}
+ ret["input_shape"] = {
+ k: v for k, v in input_shape.items() if k in cfg.MODEL.SEM_SEG_HEAD.IN_FEATURES
+ }
+ ret["conv_dim"] = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
+ ret["mask_dim"] = cfg.MODEL.SEM_SEG_HEAD.MASK_DIM
+ ret["norm"] = cfg.MODEL.SEM_SEG_HEAD.NORM
+ return ret
+
+ def forward_features(self, features):
+ multi_scale_features = []
+ num_cur_levels = 0
+ # Reverse feature maps into top-down order (from low to high resolution)
+ for idx, f in enumerate(self.in_features[::-1]):
+ x = features[f]
+ lateral_conv = self.lateral_convs[idx]
+ output_conv = self.output_convs[idx]
+ if lateral_conv is None:
+ y = output_conv(x)
+ else:
+ cur_fpn = lateral_conv(x)
+ # Following FPN implementation, we use nearest upsampling here
+ y = cur_fpn + F.interpolate(y, size=cur_fpn.shape[-2:], mode="nearest")
+ y = output_conv(y)
+ if num_cur_levels < self.maskformer_num_feature_levels:
+ multi_scale_features.append(y)
+ num_cur_levels += 1
+ return self.mask_features(y), None, multi_scale_features
+
+ def forward(self, features, targets=None):
+ logger = logging.getLogger(__name__)
+ logger.warning("Calling forward() may cause unpredicted behavior of PixelDecoder module.")
+ return self.forward_features(features)
+
+
+class TransformerEncoderOnly(nn.Module):
+ def __init__(
+ self,
+ d_model=512,
+ nhead=8,
+ num_encoder_layers=6,
+ dim_feedforward=2048,
+ dropout=0.1,
+ activation="relu",
+ normalize_before=False,
+ ):
+ super().__init__()
+
+ encoder_layer = TransformerEncoderLayer(
+ d_model, nhead, dim_feedforward, dropout, activation, normalize_before
+ )
+ encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
+ self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+ self._reset_parameters()
+
+ self.d_model = d_model
+ self.nhead = nhead
+
+ def _reset_parameters(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+
+ def forward(self, src, mask, pos_embed):
+ # flatten NxCxHxW to HWxNxC
+ bs, c, h, w = src.shape
+ src = src.flatten(2).permute(2, 0, 1)
+ pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
+ if mask is not None:
+ mask = mask.flatten(1)
+
+ memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
+ return memory.permute(1, 2, 0).view(bs, c, h, w)
+
+
+# This is a modified FPN decoder with extra Transformer encoder that processes the lowest-resolution feature map.
+@SEM_SEG_HEADS_REGISTRY.register()
+class TransformerEncoderPixelDecoder(BasePixelDecoder):
+ @configurable
+ def __init__(
+ self,
+ input_shape: Dict[str, ShapeSpec],
+ *,
+ transformer_dropout: float,
+ transformer_nheads: int,
+ transformer_dim_feedforward: int,
+ transformer_enc_layers: int,
+ transformer_pre_norm: bool,
+ conv_dim: int,
+ mask_dim: int,
+ norm: Optional[Union[str, Callable]] = None,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ input_shape: shapes (channels and stride) of the input features
+ transformer_dropout: dropout probability in transformer
+ transformer_nheads: number of heads in transformer
+ transformer_dim_feedforward: dimension of feedforward network
+ transformer_enc_layers: number of transformer encoder layers
+ transformer_pre_norm: whether to use pre-layernorm or not
+ conv_dims: number of output channels for the intermediate conv layers.
+ mask_dim: number of output channels for the final conv layer.
+ norm (str or callable): normalization for all conv layers
+ """
+ super().__init__(input_shape, conv_dim=conv_dim, mask_dim=mask_dim, norm=norm)
+
+ input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
+ self.in_features = [k for k, v in input_shape] # starting from "res2" to "res5"
+ feature_strides = [v.stride for k, v in input_shape]
+ feature_channels = [v.channels for k, v in input_shape]
+
+ in_channels = feature_channels[len(self.in_features) - 1]
+ self.input_proj = Conv2d(in_channels, conv_dim, kernel_size=1)
+ weight_init.c2_xavier_fill(self.input_proj)
+ self.transformer = TransformerEncoderOnly(
+ d_model=conv_dim,
+ dropout=transformer_dropout,
+ nhead=transformer_nheads,
+ dim_feedforward=transformer_dim_feedforward,
+ num_encoder_layers=transformer_enc_layers,
+ normalize_before=transformer_pre_norm,
+ )
+ N_steps = conv_dim // 2
+ self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
+
+ # update layer
+ use_bias = norm == ""
+ output_norm = get_norm(norm, conv_dim)
+ output_conv = Conv2d(
+ conv_dim,
+ conv_dim,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ bias=use_bias,
+ norm=output_norm,
+ activation=F.relu,
+ )
+ weight_init.c2_xavier_fill(output_conv)
+ delattr(self, "layer_{}".format(len(self.in_features)))
+ self.add_module("layer_{}".format(len(self.in_features)), output_conv)
+ self.output_convs[0] = output_conv
+
+ @classmethod
+ def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
+ ret = super().from_config(cfg, input_shape)
+ ret["transformer_dropout"] = cfg.MODEL.MASK_FORMER.DROPOUT
+ ret["transformer_nheads"] = cfg.MODEL.MASK_FORMER.NHEADS
+ ret["transformer_dim_feedforward"] = cfg.MODEL.MASK_FORMER.DIM_FEEDFORWARD
+ ret[
+ "transformer_enc_layers"
+ ] = cfg.MODEL.SEM_SEG_HEAD.TRANSFORMER_ENC_LAYERS # a separate config
+ ret["transformer_pre_norm"] = cfg.MODEL.MASK_FORMER.PRE_NORM
+ return ret
+
+ def forward_features(self, features):
+ multi_scale_features = []
+ num_cur_levels = 0
+ # Reverse feature maps into top-down order (from low to high resolution)
+ for idx, f in enumerate(self.in_features[::-1]):
+ x = features[f]
+ lateral_conv = self.lateral_convs[idx]
+ output_conv = self.output_convs[idx]
+ if lateral_conv is None:
+ transformer = self.input_proj(x)
+ pos = self.pe_layer(x)
+ transformer = self.transformer(transformer, None, pos)
+ y = output_conv(transformer)
+ # save intermediate feature as input to Transformer decoder
+ transformer_encoder_features = transformer
+ else:
+ cur_fpn = lateral_conv(x)
+ # Following FPN implementation, we use nearest upsampling here
+ y = cur_fpn + F.interpolate(y, size=cur_fpn.shape[-2:], mode="nearest")
+ y = output_conv(y)
+ if num_cur_levels < self.maskformer_num_feature_levels:
+ multi_scale_features.append(y)
+ num_cur_levels += 1
+ return self.mask_features(y), transformer_encoder_features, multi_scale_features
+
+ def forward(self, features, targets=None):
+ logger = logging.getLogger(__name__)
+ logger.warning("Calling forward() may cause unpredicted behavior of PixelDecoder module.")
+ return self.forward_features(features)

annotator/entityseg/mask2former/modeling/pixel_decoder/msdeformattn.py

@@ -0,0 +1,358 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import logging
+import numpy as np
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import fvcore.nn.weight_init as weight_init
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.init import xavier_uniform_, constant_, uniform_, normal_
+from torch.cuda.amp import autocast
+
+from detectron2.config import configurable
+from detectron2.layers import Conv2d, ShapeSpec, get_norm
+from detectron2.modeling import SEM_SEG_HEADS_REGISTRY
+
+from ..transformer_decoder.position_encoding import PositionEmbeddingSine
+from ..transformer_decoder.transformer import _get_clones, _get_activation_fn
+from .ops.modules import MSDeformAttn
+
+
+# MSDeformAttn Transformer encoder in deformable detr
+class MSDeformAttnTransformerEncoderOnly(nn.Module):
+ def __init__(self, d_model=256, nhead=8,
+ num_encoder_layers=6, dim_feedforward=1024, dropout=0.1,
+ activation="relu",
+ num_feature_levels=4, enc_n_points=4,
+ ):
+ super().__init__()
+
+ self.d_model = d_model
+ self.nhead = nhead
+
+ encoder_layer = MSDeformAttnTransformerEncoderLayer(d_model, dim_feedforward,
+ dropout, activation,
+ num_feature_levels, nhead, enc_n_points)
+ self.encoder = MSDeformAttnTransformerEncoder(encoder_layer, num_encoder_layers)
+
+ self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model))
+
+ self._reset_parameters()
+
+ def _reset_parameters(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+ for m in self.modules():
+ if isinstance(m, MSDeformAttn):
+ m._reset_parameters()
+ normal_(self.level_embed)
+
+ def get_valid_ratio(self, mask):
+ _, H, W = mask.shape
+ valid_H = torch.sum(~mask[:, :, 0], 1)
+ valid_W = torch.sum(~mask[:, 0, :], 1)
+ valid_ratio_h = valid_H.float() / H
+ valid_ratio_w = valid_W.float() / W
+ valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1)
+ return valid_ratio
+
+ def forward(self, srcs, pos_embeds):
+ masks = [torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool) for x in srcs]
+ # prepare input for encoder
+ src_flatten = []
+ mask_flatten = []
+ lvl_pos_embed_flatten = []
+ spatial_shapes = []
+ for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)):
+ bs, c, h, w = src.shape
+ spatial_shape = (h, w)
+ spatial_shapes.append(spatial_shape)
+ src = src.flatten(2).transpose(1, 2)
+ mask = mask.flatten(1)
+ pos_embed = pos_embed.flatten(2).transpose(1, 2)
+ lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)
+ lvl_pos_embed_flatten.append(lvl_pos_embed)
+ src_flatten.append(src)
+ mask_flatten.append(mask)
+ src_flatten = torch.cat(src_flatten, 1)
+ mask_flatten = torch.cat(mask_flatten, 1)
+ lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
+ spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=src_flatten.device)
+ level_start_index = torch.cat((spatial_shapes.new_zeros((1, )), spatial_shapes.prod(1).cumsum(0)[:-1]))
+ valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
+
+ # encoder
+ memory = self.encoder(src_flatten, spatial_shapes, level_start_index, valid_ratios, lvl_pos_embed_flatten, mask_flatten)
+
+ return memory, spatial_shapes, level_start_index
+
+
+class MSDeformAttnTransformerEncoderLayer(nn.Module):
+ def __init__(self,
+ d_model=256, d_ffn=1024,
+ dropout=0.1, activation="relu",
+ n_levels=4, n_heads=8, n_points=4):
+ super().__init__()
+
+ # self attention
+ self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)
+ self.dropout1 = nn.Dropout(dropout)
+ self.norm1 = nn.LayerNorm(d_model)
+
+ # ffn
+ self.linear1 = nn.Linear(d_model, d_ffn)
+ self.activation = _get_activation_fn(activation)
+ self.dropout2 = nn.Dropout(dropout)
+ self.linear2 = nn.Linear(d_ffn, d_model)
+ self.dropout3 = nn.Dropout(dropout)
+ self.norm2 = nn.LayerNorm(d_model)
+
+ @staticmethod
+ def with_pos_embed(tensor, pos):
+ return tensor if pos is None else tensor + pos
+
+ def forward_ffn(self, src):
+ src2 = self.linear2(self.dropout2(self.activation(self.linear1(src))))
+ src = src + self.dropout3(src2)
+ src = self.norm2(src)
+ return src
+
+ def forward(self, src, pos, reference_points, spatial_shapes, level_start_index, padding_mask=None):
+ # self attention
+ src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, src, spatial_shapes, level_start_index, padding_mask)
+ src = src + self.dropout1(src2)
+ src = self.norm1(src)
+
+ # ffn
+ src = self.forward_ffn(src)
+
+ return src
+
+
+class MSDeformAttnTransformerEncoder(nn.Module):
+ def __init__(self, encoder_layer, num_layers):
+ super().__init__()
+ self.layers = _get_clones(encoder_layer, num_layers)
+ self.num_layers = num_layers
+
+ @staticmethod
+ def get_reference_points(spatial_shapes, valid_ratios, device):
+ reference_points_list = []
+ for lvl, (H_, W_) in enumerate(spatial_shapes):
+
+ ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
+ torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
+ ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
+ ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
+ ref = torch.stack((ref_x, ref_y), -1)
+ reference_points_list.append(ref)
+ reference_points = torch.cat(reference_points_list, 1)
+ reference_points = reference_points[:, :, None] * valid_ratios[:, None]
+ return reference_points
+
+ def forward(self, src, spatial_shapes, level_start_index, valid_ratios, pos=None, padding_mask=None):
+ output = src
+ reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)
+ for _, layer in enumerate(self.layers):
+ output = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask)
+
+ return output
+
+
+@SEM_SEG_HEADS_REGISTRY.register()
+class MSDeformAttnPixelDecoder(nn.Module):
+ @configurable
+ def __init__(
+ self,
+ input_shape: Dict[str, ShapeSpec],
+ *,
+ transformer_dropout: float,
+ transformer_nheads: int,
+ transformer_dim_feedforward: int,
+ transformer_enc_layers: int,
+ conv_dim: int,
+ mask_dim: int,
+ norm: Optional[Union[str, Callable]] = None,
+ # deformable transformer encoder args
+ transformer_in_features: List[str],
+ common_stride: int,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ input_shape: shapes (channels and stride) of the input features
+ transformer_dropout: dropout probability in transformer
+ transformer_nheads: number of heads in transformer
+ transformer_dim_feedforward: dimension of feedforward network
+ transformer_enc_layers: number of transformer encoder layers
+ conv_dims: number of output channels for the intermediate conv layers.
+ mask_dim: number of output channels for the final conv layer.
+ norm (str or callable): normalization for all conv layers
+ """
+ super().__init__()
+ transformer_input_shape = {
+ k: v for k, v in input_shape.items() if k in transformer_in_features
+ }
+
+ # this is the input shape of pixel decoder
+ input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
+ self.in_features = [k for k, v in input_shape] # starting from "res2" to "res5"
+ self.feature_strides = [v.stride for k, v in input_shape]
+ self.feature_channels = [v.channels for k, v in input_shape]
+ 
+ # this is the input shape of transformer encoder (could use less features than pixel decoder
+ transformer_input_shape = sorted(transformer_input_shape.items(), key=lambda x: x[1].stride)
+ self.transformer_in_features = [k for k, v in transformer_input_shape] # starting from "res2" to "res5"
+ transformer_in_channels = [v.channels for k, v in transformer_input_shape]
+ self.transformer_feature_strides = [v.stride for k, v in transformer_input_shape] # to decide extra FPN layers
+
+ self.transformer_num_feature_levels = len(self.transformer_in_features)
+ if self.transformer_num_feature_levels > 1:
+ input_proj_list = []
+ # from low resolution to high resolution (res5 -> res2)
+ for in_channels in transformer_in_channels[::-1]:
+ input_proj_list.append(nn.Sequential(
+ nn.Conv2d(in_channels, conv_dim, kernel_size=1),
+ nn.GroupNorm(32, conv_dim),
+ ))
+ self.input_proj = nn.ModuleList(input_proj_list)
+ else:
+ self.input_proj = nn.ModuleList([
+ nn.Sequential(
+ nn.Conv2d(transformer_in_channels[-1], conv_dim, kernel_size=1),
+ nn.GroupNorm(32, conv_dim),
+ )])
+
+ for proj in self.input_proj:
+ nn.init.xavier_uniform_(proj[0].weight, gain=1)
+ nn.init.constant_(proj[0].bias, 0)
+
+ self.transformer = MSDeformAttnTransformerEncoderOnly(
+ d_model=conv_dim,
+ dropout=transformer_dropout,
+ nhead=transformer_nheads,
+ dim_feedforward=transformer_dim_feedforward,
+ num_encoder_layers=transformer_enc_layers,
+ num_feature_levels=self.transformer_num_feature_levels,
+ )
+ N_steps = conv_dim // 2
+ self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
+
+ self.mask_dim = mask_dim
+ # use 1x1 conv instead
+ self.mask_features = Conv2d(
+ conv_dim,
+ mask_dim,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+ weight_init.c2_xavier_fill(self.mask_features)
+ 
+ self.maskformer_num_feature_levels = 3 # always use 3 scales
+ self.common_stride = common_stride
+
+ # extra fpn levels
+ stride = min(self.transformer_feature_strides)
+ self.num_fpn_levels = int(np.log2(stride) - np.log2(self.common_stride))
+
+ lateral_convs = []
+ output_convs = []
+
+ use_bias = norm == ""
+ for idx, in_channels in enumerate(self.feature_channels[:self.num_fpn_levels]):
+ lateral_norm = get_norm(norm, conv_dim)
+ output_norm = get_norm(norm, conv_dim)
+
+ lateral_conv = Conv2d(
+ in_channels, conv_dim, kernel_size=1, bias=use_bias, norm=lateral_norm
+ )
+ output_conv = Conv2d(
+ conv_dim,
+ conv_dim,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ bias=use_bias,
+ norm=output_norm,
+ activation=F.relu,
+ )
+ weight_init.c2_xavier_fill(lateral_conv)
+ weight_init.c2_xavier_fill(output_conv)
+ self.add_module("adapter_{}".format(idx + 1), lateral_conv)
+ self.add_module("layer_{}".format(idx + 1), output_conv)
+
+ lateral_convs.append(lateral_conv)
+ output_convs.append(output_conv)
+ # Place convs into top-down order (from low to high resolution)
+ # to make the top-down computation in forward clearer.
+ self.lateral_convs = lateral_convs[::-1]
+ self.output_convs = output_convs[::-1]
+
+ @classmethod
+ def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
+ ret = {}
+ ret["input_shape"] = {
+ k: v for k, v in input_shape.items() if k in cfg.MODEL.SEM_SEG_HEAD.IN_FEATURES
+ }
+ ret["conv_dim"] = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
+ ret["mask_dim"] = cfg.MODEL.SEM_SEG_HEAD.MASK_DIM
+ ret["norm"] = cfg.MODEL.SEM_SEG_HEAD.NORM
+ ret["transformer_dropout"] = cfg.MODEL.MASK_FORMER.DROPOUT
+ ret["transformer_nheads"] = cfg.MODEL.MASK_FORMER.NHEADS
+ # ret["transformer_dim_feedforward"] = cfg.MODEL.MASK_FORMER.DIM_FEEDFORWARD
+ ret["transformer_dim_feedforward"] = 1024 # use 1024 for deformable transformer encoder
+ ret[
+ "transformer_enc_layers"
+ ] = cfg.MODEL.SEM_SEG_HEAD.TRANSFORMER_ENC_LAYERS # a separate config
+ ret["transformer_in_features"] = cfg.MODEL.SEM_SEG_HEAD.DEFORMABLE_TRANSFORMER_ENCODER_IN_FEATURES
+ ret["common_stride"] = cfg.MODEL.SEM_SEG_HEAD.COMMON_STRIDE
+ return ret
+
+ @autocast(enabled=False)
+ def forward_features(self, features):
+ srcs = []
+ pos = []
+ # Reverse feature maps into top-down order (from low to high resolution)
+ for idx, f in enumerate(self.transformer_in_features[::-1]):
+ x = features[f].float() # deformable detr does not support half precision
+ srcs.append(self.input_proj[idx](x))
+ pos.append(self.pe_layer(x))
+
+ y, spatial_shapes, level_start_index = self.transformer(srcs, pos)
+ bs = y.shape[0]
+
+ split_size_or_sections = [None] * self.transformer_num_feature_levels
+ for i in range(self.transformer_num_feature_levels):
+ if i < self.transformer_num_feature_levels - 1:
+ split_size_or_sections[i] = level_start_index[i + 1] - level_start_index[i]
+ else:
+ split_size_or_sections[i] = y.shape[1] - level_start_index[i]
+ y = torch.split(y, split_size_or_sections, dim=1)
+
+ out = []
+ multi_scale_features = []
+ num_cur_levels = 0
+ for i, z in enumerate(y):
+ out.append(z.transpose(1, 2).view(bs, -1, spatial_shapes[i][0], spatial_shapes[i][1]))
+
+ # append `out` with extra FPN levels
+ # Reverse feature maps into top-down order (from low to high resolution)
+ for idx, f in enumerate(self.in_features[:self.num_fpn_levels][::-1]):
+ x = features[f].float()
+ lateral_conv = self.lateral_convs[idx]
+ output_conv = self.output_convs[idx]
+ cur_fpn = lateral_conv(x)
+ # Following FPN implementation, we use nearest upsampling here
+ y = cur_fpn + F.interpolate(out[-1], size=cur_fpn.shape[-2:], mode="bilinear", align_corners=False)
+ y = output_conv(y)
+ out.append(y)
+
+ for o in out:
+ if num_cur_levels < self.maskformer_num_feature_levels:
+ multi_scale_features.append(o)
+ num_cur_levels += 1
+
+ return self.mask_features(out[-1]), out[0], multi_scale_features

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/functions/__init__.py

@@ -0,0 +1,13 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from .ms_deform_attn_func import MSDeformAttnFunction
+

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/functions/ms_deform_attn_func.py

@@ -0,0 +1,72 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+
+try:
+ import MultiScaleDeformableAttention as MSDA
+except ModuleNotFoundError as e:
+ info_string = (
+ "\n\nPlease compile MultiScaleDeformableAttention CUDA op with the following commands:\n"
+ "\t`cd mask2former/modeling/pixel_decoder/ops`\n"
+ "\t`sh make.sh`\n"
+ )
+ raise ModuleNotFoundError(info_string)
+
+
+class MSDeformAttnFunction(Function):
+ @staticmethod
+ def forward(ctx, value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, im2col_step):
+ ctx.im2col_step = im2col_step
+ output = MSDA.ms_deform_attn_forward(
+ value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, ctx.im2col_step)
+ ctx.save_for_backward(value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights)
+ return output
+
+ @staticmethod
+ @once_differentiable
+ def backward(ctx, grad_output):
+ value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights = ctx.saved_tensors
+ grad_value, grad_sampling_loc, grad_attn_weight = \
+ MSDA.ms_deform_attn_backward(
+ value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, grad_output, ctx.im2col_step)
+
+ return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None
+
+
+def ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights):
+ # for debug and test only,
+ # need to use cuda version instead
+ N_, S_, M_, D_ = value.shape
+ _, Lq_, M_, L_, P_, _ = sampling_locations.shape
+ value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
+ sampling_grids = 2 * sampling_locations - 1
+ sampling_value_list = []
+ for lid_, (H_, W_) in enumerate(value_spatial_shapes):
+ # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_
+ value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape(N_*M_, D_, H_, W_)
+ # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2
+ sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1)
+ # N_*M_, D_, Lq_, P_
+ sampling_value_l_ = F.grid_sample(value_l_, sampling_grid_l_,
+ mode='bilinear', padding_mode='zeros', align_corners=False)
+ sampling_value_list.append(sampling_value_l_)
+ # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_, M_, 1, Lq_, L_*P_)
+ attention_weights = attention_weights.transpose(1, 2).reshape(N_*M_, 1, Lq_, L_*P_)
+ output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(N_, M_*D_, Lq_)
+ return output.transpose(1, 2).contiguous()

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/make.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+python setup.py build install

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/modules/__init__.py

@@ -0,0 +1,12 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from .ms_deform_attn import MSDeformAttn

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/modules/ms_deform_attn.py

@@ -0,0 +1,125 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import warnings
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.init import xavier_uniform_, constant_
+
+from ..functions import MSDeformAttnFunction
+from ..functions.ms_deform_attn_func import ms_deform_attn_core_pytorch
+
+
+def _is_power_of_2(n):
+ if (not isinstance(n, int)) or (n < 0):
+ raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n)))
+ return (n & (n-1) == 0) and n != 0
+
+
+class MSDeformAttn(nn.Module):
+ def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
+ """
+ Multi-Scale Deformable Attention Module
+ :param d_model hidden dimension
+ :param n_levels number of feature levels
+ :param n_heads number of attention heads
+ :param n_points number of sampling points per attention head per feature level
+ """
+ super().__init__()
+ if d_model % n_heads != 0:
+ raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, n_heads))
+ _d_per_head = d_model // n_heads
+ # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation
+ if not _is_power_of_2(_d_per_head):
+ warnings.warn("You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 "
+ "which is more efficient in our CUDA implementation.")
+
+ self.im2col_step = 128
+
+ self.d_model = d_model
+ self.n_levels = n_levels
+ self.n_heads = n_heads
+ self.n_points = n_points
+
+ self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
+ self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
+ self.value_proj = nn.Linear(d_model, d_model)
+ self.output_proj = nn.Linear(d_model, d_model)
+
+ self._reset_parameters()
+
+ def _reset_parameters(self):
+ constant_(self.sampling_offsets.weight.data, 0.)
+ thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
+ grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+ grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1)
+ for i in range(self.n_points):
+ grid_init[:, :, i, :] *= i + 1
+ with torch.no_grad():
+ self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
+ constant_(self.attention_weights.weight.data, 0.)
+ constant_(self.attention_weights.bias.data, 0.)
+ xavier_uniform_(self.value_proj.weight.data)
+ constant_(self.value_proj.bias.data, 0.)
+ xavier_uniform_(self.output_proj.weight.data)
+ constant_(self.output_proj.bias.data, 0.)
+
+ def forward(self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, input_padding_mask=None):
+ """
+ :param query (N, Length_{query}, C)
+ :param reference_points (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area
+ or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes
+ :param input_flatten (N, \sum_{l=0}^{L-1} H_l \cdot W_l, C)
+ :param input_spatial_shapes (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]
+ :param input_level_start_index (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]
+ :param input_padding_mask (N, \sum_{l=0}^{L-1} H_l \cdot W_l), True for padding elements, False for non-padding elements
+
+ :return output (N, Length_{query}, C)
+ """
+ N, Len_q, _ = query.shape
+ N, Len_in, _ = input_flatten.shape
+ assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in
+
+ value = self.value_proj(input_flatten)
+ if input_padding_mask is not None:
+ value = value.masked_fill(input_padding_mask[..., None], float(0))
+ value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)
+ sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
+ attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
+ attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)
+ # N, Len_q, n_heads, n_levels, n_points, 2
+ if reference_points.shape[-1] == 2:
+ offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)
+ sampling_locations = reference_points[:, :, None, :, None, :] \
+ + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
+ elif reference_points.shape[-1] == 4:
+ sampling_locations = reference_points[:, :, None, :, None, :2] \
+ + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
+ else:
+ raise ValueError(
+ 'Last dim of reference_points must be 2 or 4, but get {} instead.'.format(reference_points.shape[-1]))
+ try:
+ output = MSDeformAttnFunction.apply(
+ value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)
+ except:
+ # CPU
+ output = ms_deform_attn_core_pytorch(value, input_spatial_shapes, sampling_locations, attention_weights)
+ # # For FLOPs calculation only
+ # output = ms_deform_attn_core_pytorch(value, input_spatial_shapes, sampling_locations, attention_weights)
+ output = self.output_proj(output)
+ return output

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/setup.py

@@ -0,0 +1,78 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+import os
+import glob
+
+import torch
+
+from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils.cpp_extension import CppExtension
+from torch.utils.cpp_extension import CUDAExtension
+
+from setuptools import find_packages
+from setuptools import setup
+
+requirements = ["torch", "torchvision"]
+
+def get_extensions():
+ this_dir = os.path.dirname(os.path.abspath(__file__))
+ extensions_dir = os.path.join(this_dir, "src")
+
+ main_file = glob.glob(os.path.join(extensions_dir, "*.cpp"))
+ source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp"))
+ source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))
+
+ sources = main_file + source_cpu
+ extension = CppExtension
+ extra_compile_args = {"cxx": []}
+ define_macros = []
+
+ # Force cuda since torch ask for a device, not if cuda is in fact available.
+ if (os.environ.get('FORCE_CUDA') or torch.cuda.is_available()) and CUDA_HOME is not None:
+ extension = CUDAExtension
+ sources += source_cuda
+ define_macros += [("WITH_CUDA", None)]
+ extra_compile_args["nvcc"] = [
+ "-DCUDA_HAS_FP16=1",
+ "-D__CUDA_NO_HALF_OPERATORS__",
+ "-D__CUDA_NO_HALF_CONVERSIONS__",
+ "-D__CUDA_NO_HALF2_OPERATORS__",
+ ]
+ else:
+ if CUDA_HOME is None:
+ raise NotImplementedError('CUDA_HOME is None. Please set environment variable CUDA_HOME.')
+ else:
+ raise NotImplementedError('No CUDA runtime is found. Please set FORCE_CUDA=1 or test it by running torch.cuda.is_available().')
+
+ sources = [os.path.join(extensions_dir, s) for s in sources]
+ include_dirs = [extensions_dir]
+ ext_modules = [
+ extension(
+ "MultiScaleDeformableAttention",
+ sources,
+ include_dirs=include_dirs,
+ define_macros=define_macros,
+ extra_compile_args=extra_compile_args,
+ )
+ ]
+ return ext_modules
+
+setup(
+ name="MultiScaleDeformableAttention",
+ version="1.0",
+ author="Weijie Su",
+ url="https://github.com/fundamentalvision/Deformable-DETR",
+ description="PyTorch Wrapper for CUDA Functions of Multi-Scale Deformable Attention",
+ packages=find_packages(exclude=("configs", "tests",)),
+ ext_modules=get_extensions(),
+ cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
+)

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/src/cpu/ms_deform_attn_cpu.cpp

@@ -0,0 +1,46 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+
+at::Tensor
+ms_deform_attn_cpu_forward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const int im2col_step)
+{
+ AT_ERROR("Not implement on cpu");
+}
+
+std::vector<at::Tensor>
+ms_deform_attn_cpu_backward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const at::Tensor &grad_output,
+ const int im2col_step)
+{
+ AT_ERROR("Not implement on cpu");
+}
+

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/src/cpu/ms_deform_attn_cpu.h

@@ -0,0 +1,38 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor
+ms_deform_attn_cpu_forward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const int im2col_step);
+
+std::vector<at::Tensor>
+ms_deform_attn_cpu_backward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const at::Tensor &grad_output,
+ const int im2col_step);
+
+

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/src/cuda/ms_deform_attn_cuda.cu

@@ -0,0 +1,158 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#include <vector>
+#include "cuda/ms_deform_im2col_cuda.cuh"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+
+at::Tensor ms_deform_attn_cuda_forward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const int im2col_step)
+{
+ AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+ AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
+ AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
+ AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
+ AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
+
+ AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
+ AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor");
+ AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor");
+ AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor");
+ AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor");
+
+ const int batch = value.size(0);
+ const int spatial_size = value.size(1);
+ const int num_heads = value.size(2);
+ const int channels = value.size(3);
+
+ const int num_levels = spatial_shapes.size(0);
+
+ const int num_query = sampling_loc.size(1);
+ const int num_point = sampling_loc.size(4);
+
+ const int im2col_step_ = std::min(batch, im2col_step);
+
+ AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+ 
+ auto output = at::zeros({batch, num_query, num_heads, channels}, value.options());
+
+ const int batch_n = im2col_step_;
+ auto output_n = output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels});
+ auto per_value_size = spatial_size * num_heads * channels;
+ auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
+ auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
+ for (int n = 0; n < batch/im2col_step_; ++n)
+ {
+ auto columns = output_n.select(0, n);
+ AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] {
+ ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(),
+ value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+ spatial_shapes.data<int64_t>(),
+ level_start_index.data<int64_t>(),
+ sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+ attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size,
+ batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,
+ columns.data<scalar_t>());
+
+ }));
+ }
+
+ output = output.view({batch, num_query, num_heads*channels});
+
+ return output;
+}
+
+
+std::vector<at::Tensor> ms_deform_attn_cuda_backward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const at::Tensor &grad_output,
+ const int im2col_step)
+{
+
+ AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+ AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
+ AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
+ AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
+ AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
+ AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous");
+
+ AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
+ AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor");
+ AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor");
+ AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor");
+ AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor");
+ AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor");
+
+ const int batch = value.size(0);
+ const int spatial_size = value.size(1);
+ const int num_heads = value.size(2);
+ const int channels = value.size(3);
+
+ const int num_levels = spatial_shapes.size(0);
+
+ const int num_query = sampling_loc.size(1);
+ const int num_point = sampling_loc.size(4);
+
+ const int im2col_step_ = std::min(batch, im2col_step);
+
+ AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+
+ auto grad_value = at::zeros_like(value);
+ auto grad_sampling_loc = at::zeros_like(sampling_loc);
+ auto grad_attn_weight = at::zeros_like(attn_weight);
+
+ const int batch_n = im2col_step_;
+ auto per_value_size = spatial_size * num_heads * channels;
+ auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
+ auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
+ auto grad_output_n = grad_output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels});
+ 
+ for (int n = 0; n < batch/im2col_step_; ++n)
+ {
+ auto grad_output_g = grad_output_n.select(0, n);
+ AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] {
+ ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(),
+ grad_output_g.data<scalar_t>(),
+ value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+ spatial_shapes.data<int64_t>(),
+ level_start_index.data<int64_t>(),
+ sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+ attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size,
+ batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,
+ grad_value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+ grad_sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+ grad_attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size);
+
+ }));
+ }
+
+ return {
+ grad_value, grad_sampling_loc, grad_attn_weight
+ };
+}

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/src/cuda/ms_deform_attn_cuda.h

@@ -0,0 +1,35 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor ms_deform_attn_cuda_forward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const int im2col_step);
+
+std::vector<at::Tensor> ms_deform_attn_cuda_backward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const at::Tensor &grad_output,
+ const int im2col_step);
+

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/src/cuda/ms_deform_im2col_cuda.cuh

@@ -0,0 +1,1332 @@
+/*!
+**************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************
+* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)
+* Copyright (c) 2018 Microsoft
+**************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#include <cstdio>
+#include <algorithm>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#define CUDA_KERNEL_LOOP(i, n) \
+ for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
+ i < (n); \
+ i += blockDim.x * gridDim.x)
+
+const int CUDA_NUM_THREADS = 1024;
+inline int GET_BLOCKS(const int N, const int num_threads)
+{
+ return (N + num_threads - 1) / num_threads;
+}
+
+
+template <typename scalar_t>
+__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, 
+ const int &height, const int &width, const int &nheads, const int &channels,
+ const scalar_t &h, const scalar_t &w, const int &m, const int &c)
+{
+ const int h_low = floor(h);
+ const int w_low = floor(w);
+ const int h_high = h_low + 1;
+ const int w_high = w_low + 1;
+
+ const scalar_t lh = h - h_low;
+ const scalar_t lw = w - w_low;
+ const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+ const int w_stride = nheads * channels;
+ const int h_stride = width * w_stride;
+ const int h_low_ptr_offset = h_low * h_stride;
+ const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+ const int w_low_ptr_offset = w_low * w_stride;
+ const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+ const int base_ptr = m * channels + c;
+
+ scalar_t v1 = 0;
+ if (h_low >= 0 && w_low >= 0)
+ {
+ const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+ v1 = bottom_data[ptr1];
+ }
+ scalar_t v2 = 0;
+ if (h_low >= 0 && w_high <= width - 1)
+ {
+ const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+ v2 = bottom_data[ptr2];
+ }
+ scalar_t v3 = 0;
+ if (h_high <= height - 1 && w_low >= 0)
+ {
+ const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+ v3 = bottom_data[ptr3];
+ }
+ scalar_t v4 = 0;
+ if (h_high <= height - 1 && w_high <= width - 1)
+ {
+ const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+ v4 = bottom_data[ptr4];
+ }
+
+ const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+ const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+ return val;
+}
+
+
+template <typename scalar_t>
+__device__ void ms_deform_attn_col2im_bilinear(const scalar_t* &bottom_data, 
+ const int &height, const int &width, const int &nheads, const int &channels,
+ const scalar_t &h, const scalar_t &w, const int &m, const int &c,
+ const scalar_t &top_grad,
+ const scalar_t &attn_weight,
+ scalar_t* &grad_value, 
+ scalar_t* grad_sampling_loc,
+ scalar_t* grad_attn_weight)
+{
+ const int h_low = floor(h);
+ const int w_low = floor(w);
+ const int h_high = h_low + 1;
+ const int w_high = w_low + 1;
+
+ const scalar_t lh = h - h_low;
+ const scalar_t lw = w - w_low;
+ const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+ const int w_stride = nheads * channels;
+ const int h_stride = width * w_stride;
+ const int h_low_ptr_offset = h_low * h_stride;
+ const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+ const int w_low_ptr_offset = w_low * w_stride;
+ const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+ const int base_ptr = m * channels + c;
+
+ const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+ const scalar_t top_grad_value = top_grad * attn_weight;
+ scalar_t grad_h_weight = 0, grad_w_weight = 0;
+
+ scalar_t v1 = 0;
+ if (h_low >= 0 && w_low >= 0)
+ {
+ const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+ v1 = bottom_data[ptr1];
+ grad_h_weight -= hw * v1;
+ grad_w_weight -= hh * v1;
+ atomicAdd(grad_value+ptr1, w1*top_grad_value);
+ }
+ scalar_t v2 = 0;
+ if (h_low >= 0 && w_high <= width - 1)
+ {
+ const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+ v2 = bottom_data[ptr2];
+ grad_h_weight -= lw * v2;
+ grad_w_weight += hh * v2;
+ atomicAdd(grad_value+ptr2, w2*top_grad_value);
+ }
+ scalar_t v3 = 0;
+ if (h_high <= height - 1 && w_low >= 0)
+ {
+ const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+ v3 = bottom_data[ptr3];
+ grad_h_weight += hw * v3;
+ grad_w_weight -= lh * v3;
+ atomicAdd(grad_value+ptr3, w3*top_grad_value); 
+ }
+ scalar_t v4 = 0;
+ if (h_high <= height - 1 && w_high <= width - 1)
+ {
+ const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+ v4 = bottom_data[ptr4];
+ grad_h_weight += lw * v4;
+ grad_w_weight += lh * v4;
+ atomicAdd(grad_value+ptr4, w4*top_grad_value);
+ }
+
+ const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+ *grad_attn_weight = top_grad * val;
+ *grad_sampling_loc = width * grad_w_weight * top_grad_value;
+ *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value;
+}
+
+
+template <typename scalar_t>
+__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t* &bottom_data, 
+ const int &height, const int &width, const int &nheads, const int &channels,
+ const scalar_t &h, const scalar_t &w, const int &m, const int &c,
+ const scalar_t &top_grad,
+ const scalar_t &attn_weight,
+ scalar_t* &grad_value, 
+ scalar_t* grad_sampling_loc,
+ scalar_t* grad_attn_weight)
+{
+ const int h_low = floor(h);
+ const int w_low = floor(w);
+ const int h_high = h_low + 1;
+ const int w_high = w_low + 1;
+
+ const scalar_t lh = h - h_low;
+ const scalar_t lw = w - w_low;
+ const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+ const int w_stride = nheads * channels;
+ const int h_stride = width * w_stride;
+ const int h_low_ptr_offset = h_low * h_stride;
+ const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+ const int w_low_ptr_offset = w_low * w_stride;
+ const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+ const int base_ptr = m * channels + c;
+
+ const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+ const scalar_t top_grad_value = top_grad * attn_weight;
+ scalar_t grad_h_weight = 0, grad_w_weight = 0;
+
+ scalar_t v1 = 0;
+ if (h_low >= 0 && w_low >= 0)
+ {
+ const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+ v1 = bottom_data[ptr1];
+ grad_h_weight -= hw * v1;
+ grad_w_weight -= hh * v1;
+ atomicAdd(grad_value+ptr1, w1*top_grad_value);
+ }
+ scalar_t v2 = 0;
+ if (h_low >= 0 && w_high <= width - 1)
+ {
+ const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+ v2 = bottom_data[ptr2];
+ grad_h_weight -= lw * v2;
+ grad_w_weight += hh * v2;
+ atomicAdd(grad_value+ptr2, w2*top_grad_value);
+ }
+ scalar_t v3 = 0;
+ if (h_high <= height - 1 && w_low >= 0)
+ {
+ const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+ v3 = bottom_data[ptr3];
+ grad_h_weight += hw * v3;
+ grad_w_weight -= lh * v3;
+ atomicAdd(grad_value+ptr3, w3*top_grad_value); 
+ }
+ scalar_t v4 = 0;
+ if (h_high <= height - 1 && w_high <= width - 1)
+ {
+ const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+ v4 = bottom_data[ptr4];
+ grad_h_weight += lw * v4;
+ grad_w_weight += lh * v4;
+ atomicAdd(grad_value+ptr4, w4*top_grad_value);
+ }
+
+ const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+ atomicAdd(grad_attn_weight, top_grad * val); 
+ atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value);
+ atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value);
+}
+
+
+template <typename scalar_t>
+__global__ void ms_deformable_im2col_gpu_kernel(const int n,
+ const scalar_t *data_value, 
+ const int64_t *data_spatial_shapes,
+ const int64_t *data_level_start_index, 
+ const scalar_t *data_sampling_loc,
+ const scalar_t *data_attn_weight,
+ const int batch_size, 
+ const int spatial_size, 
+ const int num_heads,
+ const int channels, 
+ const int num_levels,
+ const int num_query,
+ const int num_point,
+ scalar_t *data_col)
+{
+ CUDA_KERNEL_LOOP(index, n)
+ {
+ int _temp = index;
+ const int c_col = _temp % channels;
+ _temp /= channels;
+ const int sampling_index = _temp; 
+ const int m_col = _temp % num_heads;
+ _temp /= num_heads;
+ const int q_col = _temp % num_query;
+ _temp /= num_query;
+ const int b_col = _temp;
+
+ scalar_t *data_col_ptr = data_col + index;
+ int data_weight_ptr = sampling_index * num_levels * num_point;
+ int data_loc_w_ptr = data_weight_ptr << 1;
+ const int qid_stride = num_heads * channels;
+ const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+ scalar_t col = 0;
+ 
+ for (int l_col=0; l_col < num_levels; ++l_col)
+ {
+ const int level_start_id = data_level_start_index[l_col];
+ const int spatial_h_ptr = l_col << 1;
+ const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+ const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+ const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride);
+ for (int p_col=0; p_col < num_point; ++p_col)
+ {
+ const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+ const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+ const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+ const scalar_t h_im = loc_h * spatial_h - 0.5;
+ const scalar_t w_im = loc_w * spatial_w - 0.5;
+
+ if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+ {
+ col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight;
+ }
+
+ data_weight_ptr += 1;
+ data_loc_w_ptr += 2;
+ }
+ }
+ *data_col_ptr = col;
+ }
+}
+
+template <typename scalar_t, unsigned int blockSize>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(const int n,
+ const scalar_t *grad_col,
+ const scalar_t *data_value,
+ const int64_t *data_spatial_shapes,
+ const int64_t *data_level_start_index, 
+ const scalar_t *data_sampling_loc,
+ const scalar_t *data_attn_weight,
+ const int batch_size, 
+ const int spatial_size, 
+ const int num_heads,
+ const int channels, 
+ const int num_levels,
+ const int num_query,
+ const int num_point,
+ scalar_t *grad_value,
+ scalar_t *grad_sampling_loc,
+ scalar_t *grad_attn_weight)
+{
+ CUDA_KERNEL_LOOP(index, n)
+ {
+ __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
+ __shared__ scalar_t cache_grad_attn_weight[blockSize];
+ unsigned int tid = threadIdx.x;
+ int _temp = index;
+ const int c_col = _temp % channels;
+ _temp /= channels;
+ const int sampling_index = _temp; 
+ const int m_col = _temp % num_heads;
+ _temp /= num_heads;
+ const int q_col = _temp % num_query;
+ _temp /= num_query;
+ const int b_col = _temp;
+
+ const scalar_t top_grad = grad_col[index];
+
+ int data_weight_ptr = sampling_index * num_levels * num_point;
+ int data_loc_w_ptr = data_weight_ptr << 1;
+ const int grad_sampling_ptr = data_weight_ptr;
+ grad_sampling_loc += grad_sampling_ptr << 1;
+ grad_attn_weight += grad_sampling_ptr;
+ const int grad_weight_stride = 1;
+ const int grad_loc_stride = 2;
+ const int qid_stride = num_heads * channels;
+ const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+ for (int l_col=0; l_col < num_levels; ++l_col)
+ {
+ const int level_start_id = data_level_start_index[l_col];
+ const int spatial_h_ptr = l_col << 1;
+ const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+ const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+ const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+ const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+ scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+ for (int p_col=0; p_col < num_point; ++p_col)
+ {
+ const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+ const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+ const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+ const scalar_t h_im = loc_h * spatial_h - 0.5;
+ const scalar_t w_im = loc_w * spatial_w - 0.5;
+ *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+ *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+ *(cache_grad_attn_weight+threadIdx.x)=0;
+ if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+ {
+ ms_deform_attn_col2im_bilinear(
+ data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+ top_grad, weight, grad_value_ptr, 
+ cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+ }
+ 
+ __syncthreads();
+ if (tid == 0)
+ {
+ scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0];
+ int sid=2;
+ for (unsigned int tid = 1; tid < blockSize; ++tid)
+ {
+ _grad_w += cache_grad_sampling_loc[sid];
+ _grad_h += cache_grad_sampling_loc[sid + 1];
+ _grad_a += cache_grad_attn_weight[tid];
+ sid += 2;
+ }
+ 
+ 
+ *grad_sampling_loc = _grad_w;
+ *(grad_sampling_loc + 1) = _grad_h;
+ *grad_attn_weight = _grad_a;
+ }
+ __syncthreads();
+
+ data_weight_ptr += 1;
+ data_loc_w_ptr += 2;
+ grad_attn_weight += grad_weight_stride;
+ grad_sampling_loc += grad_loc_stride;
+ }
+ }
+ }
+}
+
+
+template <typename scalar_t, unsigned int blockSize>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(const int n,
+ const scalar_t *grad_col,
+ const scalar_t *data_value,
+ const int64_t *data_spatial_shapes,
+ const int64_t *data_level_start_index, 
+ const scalar_t *data_sampling_loc,
+ const scalar_t *data_attn_weight,
+ const int batch_size, 
+ const int spatial_size, 
+ const int num_heads,
+ const int channels, 
+ const int num_levels,
+ const int num_query,
+ const int num_point,
+ scalar_t *grad_value,
+ scalar_t *grad_sampling_loc,
+ scalar_t *grad_attn_weight)
+{
+ CUDA_KERNEL_LOOP(index, n)
+ {
+ __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
+ __shared__ scalar_t cache_grad_attn_weight[blockSize];
+ unsigned int tid = threadIdx.x;
+ int _temp = index;
+ const int c_col = _temp % channels;
+ _temp /= channels;
+ const int sampling_index = _temp; 
+ const int m_col = _temp % num_heads;
+ _temp /= num_heads;
+ const int q_col = _temp % num_query;
+ _temp /= num_query;
+ const int b_col = _temp;
+
+ const scalar_t top_grad = grad_col[index];
+
+ int data_weight_ptr = sampling_index * num_levels * num_point;
+ int data_loc_w_ptr = data_weight_ptr << 1;
+ const int grad_sampling_ptr = data_weight_ptr;
+ grad_sampling_loc += grad_sampling_ptr << 1;
+ grad_attn_weight += grad_sampling_ptr;
+ const int grad_weight_stride = 1;
+ const int grad_loc_stride = 2;
+ const int qid_stride = num_heads * channels;
+ const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+ for (int l_col=0; l_col < num_levels; ++l_col)
+ {
+ const int level_start_id = data_level_start_index[l_col];
+ const int spatial_h_ptr = l_col << 1;
+ const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+ const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+ const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+ const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+ scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+ for (int p_col=0; p_col < num_point; ++p_col)
+ {
+ const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+ const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+ const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+ const scalar_t h_im = loc_h * spatial_h - 0.5;
+ const scalar_t w_im = loc_w * spatial_w - 0.5;
+ *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+ *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+ *(cache_grad_attn_weight+threadIdx.x)=0;
+ if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+ {
+ ms_deform_attn_col2im_bilinear(
+ data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+ top_grad, weight, grad_value_ptr, 
+ cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+ }
+ 
+ __syncthreads();
+
+ for (unsigned int s=blockSize/2; s>0; s>>=1)
+ {
+ if (tid < s) {
+ const unsigned int xid1 = tid << 1;
+ const unsigned int xid2 = (tid + s) << 1;
+ cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+ cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+ cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+ }
+ __syncthreads();
+ }
+
+ if (tid == 0)
+ { 
+ *grad_sampling_loc = cache_grad_sampling_loc[0];
+ *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
+ *grad_attn_weight = cache_grad_attn_weight[0];
+ }
+ __syncthreads();
+
+ data_weight_ptr += 1;
+ data_loc_w_ptr += 2;
+ grad_attn_weight += grad_weight_stride;
+ grad_sampling_loc += grad_loc_stride;
+ }
+ }
+ }
+}
+
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(const int n,
+ const scalar_t *grad_col,
+ const scalar_t *data_value,
+ const int64_t *data_spatial_shapes,
+ const int64_t *data_level_start_index, 
+ const scalar_t *data_sampling_loc,
+ const scalar_t *data_attn_weight,
+ const int batch_size, 
+ const int spatial_size, 
+ const int num_heads,
+ const int channels, 
+ const int num_levels,
+ const int num_query,
+ const int num_point,
+ scalar_t *grad_value,
+ scalar_t *grad_sampling_loc,
+ scalar_t *grad_attn_weight)
+{
+ CUDA_KERNEL_LOOP(index, n)
+ {
+ extern __shared__ int _s[];
+ scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;
+ scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+ unsigned int tid = threadIdx.x;
+ int _temp = index;
+ const int c_col = _temp % channels;
+ _temp /= channels;
+ const int sampling_index = _temp; 
+ const int m_col = _temp % num_heads;
+ _temp /= num_heads;
+ const int q_col = _temp % num_query;
+ _temp /= num_query;
+ const int b_col = _temp;
+
+ const scalar_t top_grad = grad_col[index];
+
+ int data_weight_ptr = sampling_index * num_levels * num_point;
+ int data_loc_w_ptr = data_weight_ptr << 1;
+ const int grad_sampling_ptr = data_weight_ptr;
+ grad_sampling_loc += grad_sampling_ptr << 1;
+ grad_attn_weight += grad_sampling_ptr;
+ const int grad_weight_stride = 1;
+ const int grad_loc_stride = 2;
+ const int qid_stride = num_heads * channels;
+ const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+ for (int l_col=0; l_col < num_levels; ++l_col)
+ {
+ const int level_start_id = data_level_start_index[l_col];
+ const int spatial_h_ptr = l_col << 1;
+ const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+ const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+ const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+ const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+ scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+ for (int p_col=0; p_col < num_point; ++p_col)
+ {
+ const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+ const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+ const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+ const scalar_t h_im = loc_h * spatial_h - 0.5;
+ const scalar_t w_im = loc_w * spatial_w - 0.5;
+ *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+ *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+ *(cache_grad_attn_weight+threadIdx.x)=0;
+ if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+ {
+ ms_deform_attn_col2im_bilinear(
+ data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+ top_grad, weight, grad_value_ptr, 
+ cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+ }
+ 
+ __syncthreads();
+ if (tid == 0)
+ {
+ scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0];
+ int sid=2;
+ for (unsigned int tid = 1; tid < blockDim.x; ++tid)
+ {
+ _grad_w += cache_grad_sampling_loc[sid];
+ _grad_h += cache_grad_sampling_loc[sid + 1];
+ _grad_a += cache_grad_attn_weight[tid];
+ sid += 2;
+ }
+ 
+ 
+ *grad_sampling_loc = _grad_w;
+ *(grad_sampling_loc + 1) = _grad_h;
+ *grad_attn_weight = _grad_a;
+ }
+ __syncthreads();
+
+ data_weight_ptr += 1;
+ data_loc_w_ptr += 2;
+ grad_attn_weight += grad_weight_stride;
+ grad_sampling_loc += grad_loc_stride;
+ }
+ }
+ }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(const int n,
+ const scalar_t *grad_col,
+ const scalar_t *data_value,
+ const int64_t *data_spatial_shapes,
+ const int64_t *data_level_start_index, 
+ const scalar_t *data_sampling_loc,
+ const scalar_t *data_attn_weight,
+ const int batch_size, 
+ const int spatial_size, 
+ const int num_heads,
+ const int channels, 
+ const int num_levels,
+ const int num_query,
+ const int num_point,
+ scalar_t *grad_value,
+ scalar_t *grad_sampling_loc,
+ scalar_t *grad_attn_weight)
+{
+ CUDA_KERNEL_LOOP(index, n)
+ {
+ extern __shared__ int _s[];
+ scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;
+ scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+ unsigned int tid = threadIdx.x;
+ int _temp = index;
+ const int c_col = _temp % channels;
+ _temp /= channels;
+ const int sampling_index = _temp; 
+ const int m_col = _temp % num_heads;
+ _temp /= num_heads;
+ const int q_col = _temp % num_query;
+ _temp /= num_query;
+ const int b_col = _temp;
+
+ const scalar_t top_grad = grad_col[index];
+
+ int data_weight_ptr = sampling_index * num_levels * num_point;
+ int data_loc_w_ptr = data_weight_ptr << 1;
+ const int grad_sampling_ptr = data_weight_ptr;
+ grad_sampling_loc += grad_sampling_ptr << 1;
+ grad_attn_weight += grad_sampling_ptr;
+ const int grad_weight_stride = 1;
+ const int grad_loc_stride = 2;
+ const int qid_stride = num_heads * channels;
+ const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+ for (int l_col=0; l_col < num_levels; ++l_col)
+ {
+ const int level_start_id = data_level_start_index[l_col];
+ const int spatial_h_ptr = l_col << 1;
+ const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+ const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+ const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+ const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+ scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+ for (int p_col=0; p_col < num_point; ++p_col)
+ {
+ const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+ const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+ const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+ const scalar_t h_im = loc_h * spatial_h - 0.5;
+ const scalar_t w_im = loc_w * spatial_w - 0.5;
+ *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+ *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+ *(cache_grad_attn_weight+threadIdx.x)=0;
+ if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+ {
+ ms_deform_attn_col2im_bilinear(
+ data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+ top_grad, weight, grad_value_ptr, 
+ cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+ }
+ 
+ __syncthreads();
+
+ for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1)
+ {
+ if (tid < s) {
+ const unsigned int xid1 = tid << 1;
+ const unsigned int xid2 = (tid + s) << 1;
+ cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+ cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+ cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+ if (tid + (s << 1) < spre)
+ {
+ cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
+ cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
+ cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
+ } 
+ }
+ __syncthreads();
+ }
+
+ if (tid == 0)
+ {
+ *grad_sampling_loc = cache_grad_sampling_loc[0];
+ *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
+ *grad_attn_weight = cache_grad_attn_weight[0];
+ }
+ __syncthreads();
+
+ data_weight_ptr += 1;
+ data_loc_w_ptr += 2;
+ grad_attn_weight += grad_weight_stride;
+ grad_sampling_loc += grad_loc_stride;
+ }
+ }
+ }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(const int n,
+ const scalar_t *grad_col,
+ const scalar_t *data_value,
+ const int64_t *data_spatial_shapes,
+ const int64_t *data_level_start_index, 
+ const scalar_t *data_sampling_loc,
+ const scalar_t *data_attn_weight,
+ const int batch_size, 
+ const int spatial_size, 
+ const int num_heads,
+ const int channels, 
+ const int num_levels,
+ const int num_query,
+ const int num_point,
+ scalar_t *grad_value,
+ scalar_t *grad_sampling_loc,
+ scalar_t *grad_attn_weight)
+{
+ CUDA_KERNEL_LOOP(index, n)
+ {
+ extern __shared__ int _s[];
+ scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;
+ scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+ unsigned int tid = threadIdx.x;
+ int _temp = index;
+ const int c_col = _temp % channels;
+ _temp /= channels;
+ const int sampling_index = _temp; 
+ const int m_col = _temp % num_heads;
+ _temp /= num_heads;
+ const int q_col = _temp % num_query;
+ _temp /= num_query;
+ const int b_col = _temp;
+
+ const scalar_t top_grad = grad_col[index];
+
+ int data_weight_ptr = sampling_index * num_levels * num_point;
+ int data_loc_w_ptr = data_weight_ptr << 1;
+ const int grad_sampling_ptr = data_weight_ptr;
+ grad_sampling_loc += grad_sampling_ptr << 1;
+ grad_attn_weight += grad_sampling_ptr;
+ const int grad_weight_stride = 1;
+ const int grad_loc_stride = 2;
+ const int qid_stride = num_heads * channels;
+ const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+ for (int l_col=0; l_col < num_levels; ++l_col)
+ {
+ const int level_start_id = data_level_start_index[l_col];
+ const int spatial_h_ptr = l_col << 1;
+ const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+ const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+ const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+ const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+ scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+ for (int p_col=0; p_col < num_point; ++p_col)
+ {
+ const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+ const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+ const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+ const scalar_t h_im = loc_h * spatial_h - 0.5;
+ const scalar_t w_im = loc_w * spatial_w - 0.5;
+ *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+ *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+ *(cache_grad_attn_weight+threadIdx.x)=0;
+ if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+ {
+ ms_deform_attn_col2im_bilinear(
+ data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+ top_grad, weight, grad_value_ptr, 
+ cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+ }
+ 
+ __syncthreads();
+
+ for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1)
+ {
+ if (tid < s) {
+ const unsigned int xid1 = tid << 1;
+ const unsigned int xid2 = (tid + s) << 1;
+ cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+ cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+ cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+ if (tid + (s << 1) < spre)
+ {
+ cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
+ cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
+ cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
+ }
+ }
+ __syncthreads();
+ }
+
+ if (tid == 0)
+ {
+ atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]);
+ atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]);
+ atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]);
+ }
+ __syncthreads();
+
+ data_weight_ptr += 1;
+ data_loc_w_ptr += 2;
+ grad_attn_weight += grad_weight_stride;
+ grad_sampling_loc += grad_loc_stride;
+ }
+ }
+ }
+}
+
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n,
+ const scalar_t *grad_col,
+ const scalar_t *data_value,
+ const int64_t *data_spatial_shapes,
+ const int64_t *data_level_start_index, 
+ const scalar_t *data_sampling_loc,
+ const scalar_t *data_attn_weight,
+ const int batch_size, 
+ const int spatial_size, 
+ const int num_heads,
+ const int channels, 
+ const int num_levels,
+ const int num_query,
+ const int num_point,
+ scalar_t *grad_value,
+ scalar_t *grad_sampling_loc,
+ scalar_t *grad_attn_weight)
+{
+ CUDA_KERNEL_LOOP(index, n)
+ {
+ int _temp = index;
+ const int c_col = _temp % channels;
+ _temp /= channels;
+ const int sampling_index = _temp; 
+ const int m_col = _temp % num_heads;
+ _temp /= num_heads;
+ const int q_col = _temp % num_query;
+ _temp /= num_query;
+ const int b_col = _temp;
+
+ const scalar_t top_grad = grad_col[index];
+
+ int data_weight_ptr = sampling_index * num_levels * num_point;
+ int data_loc_w_ptr = data_weight_ptr << 1;
+ const int grad_sampling_ptr = data_weight_ptr;
+ grad_sampling_loc += grad_sampling_ptr << 1;
+ grad_attn_weight += grad_sampling_ptr;
+ const int grad_weight_stride = 1;
+ const int grad_loc_stride = 2;
+ const int qid_stride = num_heads * channels;
+ const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+ for (int l_col=0; l_col < num_levels; ++l_col)
+ {
+ const int level_start_id = data_level_start_index[l_col];
+ const int spatial_h_ptr = l_col << 1;
+ const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+ const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+ const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+ const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+ scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+ for (int p_col=0; p_col < num_point; ++p_col)
+ {
+ const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+ const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+ const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+ const scalar_t h_im = loc_h * spatial_h - 0.5;
+ const scalar_t w_im = loc_w * spatial_w - 0.5;
+ if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+ {
+ ms_deform_attn_col2im_bilinear_gm(
+ data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+ top_grad, weight, grad_value_ptr, 
+ grad_sampling_loc, grad_attn_weight);
+ }
+ data_weight_ptr += 1;
+ data_loc_w_ptr += 2;
+ grad_attn_weight += grad_weight_stride;
+ grad_sampling_loc += grad_loc_stride;
+ }
+ }
+ }
+}
+
+
+template <typename scalar_t>
+void ms_deformable_im2col_cuda(cudaStream_t stream,
+ const scalar_t* data_value,
+ const int64_t* data_spatial_shapes, 
+ const int64_t* data_level_start_index, 
+ const scalar_t* data_sampling_loc,
+ const scalar_t* data_attn_weight,
+ const int batch_size,
+ const int spatial_size, 
+ const int num_heads, 
+ const int channels, 
+ const int num_levels, 
+ const int num_query,
+ const int num_point,
+ scalar_t* data_col)
+{
+ const int num_kernels = batch_size * num_query * num_heads * channels;
+ const int num_actual_kernels = batch_size * num_query * num_heads * channels;
+ const int num_threads = CUDA_NUM_THREADS;
+ ms_deformable_im2col_gpu_kernel<scalar_t>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, 
+ batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col);
+ 
+ cudaError_t err = cudaGetLastError();
+ if (err != cudaSuccess)
+ {
+ printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err));
+ }
+
+}
+
+template <typename scalar_t>
+void ms_deformable_col2im_cuda(cudaStream_t stream,
+ const scalar_t* grad_col,
+ const scalar_t* data_value,
+ const int64_t * data_spatial_shapes,
+ const int64_t * data_level_start_index,
+ const scalar_t * data_sampling_loc,
+ const scalar_t * data_attn_weight,
+ const int batch_size, 
+ const int spatial_size, 
+ const int num_heads,
+ const int channels, 
+ const int num_levels,
+ const int num_query,
+ const int num_point, 
+ scalar_t* grad_value,
+ scalar_t* grad_sampling_loc,
+ scalar_t* grad_attn_weight)
+{
+ const int num_threads = (channels > CUDA_NUM_THREADS)?CUDA_NUM_THREADS:channels;
+ const int num_kernels = batch_size * num_query * num_heads * channels;
+ const int num_actual_kernels = batch_size * num_query * num_heads * channels;
+ if (channels > 1024)
+ {
+ if ((channels & 1023) == 0)
+ {
+ ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks<scalar_t>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ num_threads*3*sizeof(scalar_t), stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ }
+ else
+ {
+ ms_deformable_col2im_gpu_kernel_gm<scalar_t>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ }
+ }
+ else{
+ switch(channels)
+ {
+ case 1:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 1>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 2:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 2>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 4:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 4>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 8:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 8>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 16:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 16>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 32:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 32>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 64:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 64>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 128:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 128>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 256:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 256>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 512:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 512>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ case 1024:
+ ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 1024>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ 0, stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ break;
+ default:
+ if (channels < 64)
+ {
+ ms_deformable_col2im_gpu_kernel_shm_reduce_v1<scalar_t>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ num_threads*3*sizeof(scalar_t), stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ }
+ else
+ {
+ ms_deformable_col2im_gpu_kernel_shm_reduce_v2<scalar_t>
+ <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+ num_threads*3*sizeof(scalar_t), stream>>>(
+ num_kernels, 
+ grad_col,
+ data_value,
+ data_spatial_shapes,
+ data_level_start_index, 
+ data_sampling_loc,
+ data_attn_weight,
+ batch_size, 
+ spatial_size, 
+ num_heads,
+ channels, 
+ num_levels,
+ num_query,
+ num_point,
+ grad_value,
+ grad_sampling_loc,
+ grad_attn_weight);
+ }
+ }
+ }
+ cudaError_t err = cudaGetLastError();
+ if (err != cudaSuccess)
+ {
+ printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err));
+ }
+
+}

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/src/ms_deform_attn.h

@@ -0,0 +1,67 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#pragma once
+
+#include "cpu/ms_deform_attn_cpu.h"
+
+#ifdef WITH_CUDA
+#include "cuda/ms_deform_attn_cuda.h"
+#endif
+
+
+at::Tensor
+ms_deform_attn_forward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const int im2col_step)
+{
+ if (value.type().is_cuda())
+ {
+#ifdef WITH_CUDA
+ return ms_deform_attn_cuda_forward(
+ value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step);
+#else
+ AT_ERROR("Not compiled with GPU support");
+#endif
+ }
+ AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor>
+ms_deform_attn_backward(
+ const at::Tensor &value, 
+ const at::Tensor &spatial_shapes,
+ const at::Tensor &level_start_index,
+ const at::Tensor &sampling_loc,
+ const at::Tensor &attn_weight,
+ const at::Tensor &grad_output,
+ const int im2col_step)
+{
+ if (value.type().is_cuda())
+ {
+#ifdef WITH_CUDA
+ return ms_deform_attn_cuda_backward(
+ value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step);
+#else
+ AT_ERROR("Not compiled with GPU support");
+#endif
+ }
+ AT_ERROR("Not implemented on the CPU");
+}
+

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/src/vision.cpp

@@ -0,0 +1,21 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#include "ms_deform_attn.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("ms_deform_attn_forward", &ms_deform_attn_forward, "ms_deform_attn_forward");
+ m.def("ms_deform_attn_backward", &ms_deform_attn_backward, "ms_deform_attn_backward");
+}

annotator/entityseg/mask2former/modeling/pixel_decoder/ops/test.py

@@ -0,0 +1,92 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+from torch.autograd import gradcheck
+
+from functions.ms_deform_attn_func import MSDeformAttnFunction, ms_deform_attn_core_pytorch
+
+
+N, M, D = 1, 2, 2
+Lq, L, P = 2, 2, 2
+shapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()
+level_start_index = torch.cat((shapes.new_zeros((1, )), shapes.prod(1).cumsum(0)[:-1]))
+S = sum([(H*W).item() for H, W in shapes])
+
+
+torch.manual_seed(3)
+
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_double():
+ value = torch.rand(N, S, M, D).cuda() * 0.01
+ sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+ attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+ attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
+ im2col_step = 2
+ output_pytorch = ms_deform_attn_core_pytorch(value.double(), shapes, sampling_locations.double(), attention_weights.double()).detach().cpu()
+ output_cuda = MSDeformAttnFunction.apply(value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step).detach().cpu()
+ fwdok = torch.allclose(output_cuda, output_pytorch)
+ max_abs_err = (output_cuda - output_pytorch).abs().max()
+ max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
+
+ print(f'* {fwdok} check_forward_equal_with_pytorch_double: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_float():
+ value = torch.rand(N, S, M, D).cuda() * 0.01
+ sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+ attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+ attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
+ im2col_step = 2
+ output_pytorch = ms_deform_attn_core_pytorch(value, shapes, sampling_locations, attention_weights).detach().cpu()
+ output_cuda = MSDeformAttnFunction.apply(value, shapes, level_start_index, sampling_locations, attention_weights, im2col_step).detach().cpu()
+ fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+ max_abs_err = (output_cuda - output_pytorch).abs().max()
+ max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
+
+ print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+
+def check_gradient_numerical(channels=4, grad_value=True, grad_sampling_loc=True, grad_attn_weight=True):
+
+ value = torch.rand(N, S, M, channels).cuda() * 0.01
+ sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+ attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+ attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
+ im2col_step = 2
+ func = MSDeformAttnFunction.apply
+
+ value.requires_grad = grad_value
+ sampling_locations.requires_grad = grad_sampling_loc
+ attention_weights.requires_grad = grad_attn_weight
+
+ gradok = gradcheck(func, (value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step))
+
+ print(f'* {gradok} check_gradient_numerical(D={channels})')
+
+
+if __name__ == '__main__':
+ check_forward_equal_with_pytorch_double()
+ check_forward_equal_with_pytorch_float()
+
+ for channels in [30, 32, 64, 71, 1025, 2048, 3096]:
+ check_gradient_numerical(channels, True, True, True)
+
+
+

annotator/entityseg/mask2former/modeling/transformer_decoder/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from .maskformer_transformer_decoder import StandardTransformerDecoder
+from .mask2former_transformer_decoder import MultiScaleMaskedTransformerDecoder
+from .cropformer_transformer_decoder import CropSharedMultiScaleMaskedTransformerDecoder
+

annotator/entityseg/mask2former/modeling/transformer_decoder/cropformer_transformer_decoder.py

@@ -0,0 +1,595 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/detr.py
+import logging
+import fvcore.nn.weight_init as weight_init
+from typing import Optional
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+from detectron2.config import configurable
+from detectron2.layers import Conv2d
+
+from .position_encoding import PositionEmbeddingSine3D2D
+from .maskformer_transformer_decoder import TRANSFORMER_DECODER_REGISTRY
+
+import pdb
+
+class SelfAttentionLayer(nn.Module):
+
+ def __init__(self, d_model, nhead, dropout=0.0,
+ activation="relu", normalize_before=False):
+ super().__init__()
+ self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+
+ self.norm = nn.LayerNorm(d_model)
+ self.dropout = nn.Dropout(dropout)
+
+ self.activation = _get_activation_fn(activation)
+ self.normalize_before = normalize_before
+
+ self._reset_parameters()
+ 
+ def _reset_parameters(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+
+ def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+ return tensor if pos is None else tensor + pos
+
+ def forward_post(self, tgt,
+ tgt_mask: Optional[Tensor] = None,
+ tgt_key_padding_mask: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ q = k = self.with_pos_embed(tgt, query_pos)
+ tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
+ key_padding_mask=tgt_key_padding_mask)[0]
+ tgt = tgt + self.dropout(tgt2)
+ tgt = self.norm(tgt)
+
+ return tgt
+
+ def forward_pre(self, tgt,
+ tgt_mask: Optional[Tensor] = None,
+ tgt_key_padding_mask: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ tgt2 = self.norm(tgt)
+ q = k = self.with_pos_embed(tgt2, query_pos)
+ tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+ key_padding_mask=tgt_key_padding_mask)[0]
+ tgt = tgt + self.dropout(tgt2)
+ 
+ return tgt
+
+ def forward(self, tgt,
+ tgt_mask: Optional[Tensor] = None,
+ tgt_key_padding_mask: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ if self.normalize_before:
+ return self.forward_pre(tgt, tgt_mask,
+ tgt_key_padding_mask, query_pos)
+ return self.forward_post(tgt, tgt_mask,
+ tgt_key_padding_mask, query_pos)
+
+
+class CrossAttentionLayer(nn.Module):
+
+ def __init__(self, d_model, nhead, dropout=0.0,
+ activation="relu", normalize_before=False):
+ super().__init__()
+ self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+
+ self.norm = nn.LayerNorm(d_model)
+ self.dropout = nn.Dropout(dropout)
+
+ self.activation = _get_activation_fn(activation)
+ self.normalize_before = normalize_before
+
+ self._reset_parameters()
+ 
+ def _reset_parameters(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+
+ def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+ return tensor if pos is None else tensor + pos
+
+ def forward_post(self, tgt, memory,
+ memory_mask: Optional[Tensor] = None,
+ memory_key_padding_mask: Optional[Tensor] = None,
+ pos: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
+ key=self.with_pos_embed(memory, pos),
+ value=memory, attn_mask=memory_mask,
+ key_padding_mask=memory_key_padding_mask)[0]
+ tgt = tgt + self.dropout(tgt2)
+ tgt = self.norm(tgt)
+ 
+ return tgt
+
+ def forward_pre(self, tgt, memory,
+ memory_mask: Optional[Tensor] = None,
+ memory_key_padding_mask: Optional[Tensor] = None,
+ pos: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ tgt2 = self.norm(tgt)
+ tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+ key=self.with_pos_embed(memory, pos),
+ value=memory, attn_mask=memory_mask,
+ key_padding_mask=memory_key_padding_mask)[0]
+ tgt = tgt + self.dropout(tgt2)
+
+ return tgt
+
+ def forward(self, tgt, memory,
+ memory_mask: Optional[Tensor] = None,
+ memory_key_padding_mask: Optional[Tensor] = None,
+ pos: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ if self.normalize_before:
+ return self.forward_pre(tgt, memory, memory_mask,
+ memory_key_padding_mask, pos, query_pos)
+ return self.forward_post(tgt, memory, memory_mask,
+ memory_key_padding_mask, pos, query_pos)
+
+
+class FFNLayer(nn.Module):
+
+ def __init__(self, d_model, dim_feedforward=2048, dropout=0.0,
+ activation="relu", normalize_before=False):
+ super().__init__()
+ # Implementation of Feedforward model
+ self.linear1 = nn.Linear(d_model, dim_feedforward)
+ self.dropout = nn.Dropout(dropout)
+ self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+ self.norm = nn.LayerNorm(d_model)
+
+ self.activation = _get_activation_fn(activation)
+ self.normalize_before = normalize_before
+
+ self._reset_parameters()
+ 
+ def _reset_parameters(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+
+ def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+ return tensor if pos is None else tensor + pos
+
+ def forward_post(self, tgt):
+ tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+ tgt = tgt + self.dropout(tgt2)
+ tgt = self.norm(tgt)
+ return tgt
+
+ def forward_pre(self, tgt):
+ tgt2 = self.norm(tgt)
+ tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+ tgt = tgt + self.dropout(tgt2)
+ return tgt
+
+ def forward(self, tgt):
+ if self.normalize_before:
+ return self.forward_pre(tgt)
+ return self.forward_post(tgt)
+
+def _get_activation_fn(activation):
+ """Return an activation function given a string"""
+ if activation == "relu":
+ return F.relu
+ if activation == "gelu":
+ return F.gelu
+ if activation == "glu":
+ return F.glu
+ raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
+
+
+class MLP(nn.Module):
+ """ Very simple multi-layer perceptron (also called FFN)"""
+
+ def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+ super().__init__()
+ self.num_layers = num_layers
+ h = [hidden_dim] * (num_layers - 1)
+ self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+ def forward(self, x):
+ for i, layer in enumerate(self.layers):
+ x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+ return x
+
+class Make3dQueries(nn.Module):
+ _version = 2
+ def __init__(self, cfg):
+ super().__init__()
+ self.cfg = cfg
+ self.enc_crosattn_3d = nn.ModuleList()
+ self.enc_selfattn_3d = nn.ModuleList()
+ self.enc_ffn_3d = nn.ModuleList()
+ self.num_layers_3d = cfg.ENTITY.FUSE_NUM_LAYERS
+ for _ in range(self.num_layers_3d):
+ self.enc_crosattn_3d.append(
+ CrossAttentionLayer(
+ d_model=cfg.ENTITY.FUSE_ENC_HIDDIEN_DIM,
+ nhead=cfg.ENTITY.FUSE_ENC_NHEADS,
+ dropout=0.0,
+ normalize_before=cfg.ENTITY.FUSE_ENC_PRE_NORM)
+ )
+ self.enc_selfattn_3d.append(
+ SelfAttentionLayer(
+ d_model=cfg.ENTITY.FUSE_ENC_HIDDIEN_DIM,
+ nhead=cfg.ENTITY.FUSE_ENC_NHEADS,
+ dropout=0.0,
+ normalize_before=cfg.ENTITY.FUSE_ENC_PRE_NORM)
+ )
+ self.enc_ffn_3d.append(
+ FFNLayer(
+ d_model=cfg.ENTITY.FUSE_ENC_HIDDIEN_DIM,
+ dim_feedforward=cfg.ENTITY.FUSE_ENC_DIM_FEEDFORWARD,
+ dropout=0.0,
+ normalize_before=cfg.ENTITY.FUSE_ENC_PRE_NORM,
+ )
+ )
+ 
+ def forward(self, output_2d, query_embed_2d, query_embed_3d):
+ Q, BT, C = query_embed_2d.shape
+ Q, B, C = query_embed_3d.shape
+ T = int(BT / B)
+
+ output_3d = output_2d[:,0::T,:]
+ ### (Q, B, T, C)
+ output_2d = output_2d.unflatten(1, (B, T)).permute((0,2,1,3)).flatten(0,1)
+ query_embed_2d = query_embed_2d.unflatten(1, (B, T)).permute((0,2,1,3)).flatten(0,1)
+
+ for i in range(self.num_layers_3d):
+ output_3d = self.enc_crosattn_3d[i](output_3d, output_2d, pos=query_embed_2d, query_pos=query_embed_3d)
+ output_3d = self.enc_selfattn_3d[i](output_3d)
+ output_3d = self.enc_ffn_3d[i](output_3d)
+ 
+ return output_3d
+
+
+@TRANSFORMER_DECODER_REGISTRY.register()
+class CropSharedMultiScaleMaskedTransformerDecoder(nn.Module):
+ _version = 2
+
+ def _load_from_state_dict(
+ self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+ ):
+ version = local_metadata.get("version", None)
+ if version is None or version < 2:
+ # Do not warn if train from scratch
+ scratch = True
+ logger = logging.getLogger(__name__)
+ for k in list(state_dict.keys()):
+ newk = k
+ if "static_query" in k:
+ newk = k.replace("static_query", "query_feat")
+ if newk != k:
+ state_dict[newk] = state_dict[k]
+ del state_dict[k]
+ scratch = False
+
+ if not scratch:
+ logger.warning(
+ f"Weight format of {self.__class__.__name__} have changed! "
+ "Please upgrade your models. Applying automatic conversion now ..."
+ )
+
+ @configurable
+ def __init__(
+ self,
+ cfg,
+ in_channels,
+ mask_classification=True,
+ *,
+ num_classes: int,
+ hidden_dim: int,
+ num_queries: int,
+ nheads: int,
+ dim_feedforward: int,
+ dec_layers: int,
+ pre_norm: bool,
+ mask_dim: int,
+ enforce_input_project: bool,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ in_channels: channels of the input features
+ mask_classification: whether to add mask classifier or not
+ num_classes: number of classes
+ hidden_dim: Transformer feature dimension
+ num_queries: number of queries
+ nheads: number of heads
+ dim_feedforward: feature dimension in feedforward network
+ enc_layers: number of Transformer encoder layers
+ dec_layers: number of Transformer decoder layers
+ pre_norm: whether to use pre-LayerNorm or not
+ mask_dim: mask feature dimension
+ enforce_input_project: add input project 1x1 conv even if input
+ channels and hidden dim is identical
+ """
+ super().__init__()
+
+ assert mask_classification, "Only support mask classification model"
+ self.cfg = cfg
+
+ self.mask_classification = mask_classification
+ # positional encoding
+ N_steps = hidden_dim // 2
+ 
+ self.pe_layer = PositionEmbeddingSine3D2D(N_steps, normalize=True)
+ 
+ # define Transformer decoder here
+ self.num_heads = nheads
+ self.num_layers = dec_layers
+ self.transformer_self_attention_layers = nn.ModuleList()
+ self.transformer_cross_attention_layers = nn.ModuleList()
+ self.transformer_ffn_layers = nn.ModuleList()
+
+ for _ in range(self.num_layers):
+ self.transformer_self_attention_layers.append(
+ SelfAttentionLayer(
+ d_model=hidden_dim,
+ nhead=nheads,
+ dropout=0.0,
+ normalize_before=pre_norm,
+ )
+ )
+
+ self.transformer_cross_attention_layers.append(
+ CrossAttentionLayer(
+ d_model=hidden_dim,
+ nhead=nheads,
+ dropout=0.0,
+ normalize_before=pre_norm,
+ )
+ )
+
+ self.transformer_ffn_layers.append(
+ FFNLayer(
+ d_model=hidden_dim,
+ dim_feedforward=dim_feedforward,
+ dropout=0.0,
+ normalize_before=pre_norm,
+ )
+ )
+
+ self.make_3d = Make3dQueries(cfg)
+ self.decoder_norm = nn.LayerNorm(hidden_dim)
+
+ self.num_queries = num_queries
+ # learnable query features
+ self.query_feat = nn.Embedding(num_queries, hidden_dim)
+ # learnable query p.e.
+ self.query_embed = nn.Embedding(num_queries, hidden_dim)
+
+ # level embedding (we always use 3 scales)
+ self.num_feature_levels = 3
+ self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim)
+ self.input_proj = nn.ModuleList()
+ for _ in range(self.num_feature_levels):
+ if in_channels != hidden_dim or enforce_input_project:
+ self.input_proj.append(Conv2d(in_channels, hidden_dim, kernel_size=1))
+ weight_init.c2_xavier_fill(self.input_proj[-1])
+ weight_init.c2_xavier_fill(self.input_proj_3d[-1])
+ else:
+ self.input_proj.append(nn.Sequential())
+
+ # output FFNs
+ if self.mask_classification:
+ self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
+ self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3)
+
+ @classmethod
+ def from_config(cls, cfg, in_channels, mask_classification):
+ ret = {}
+ ret["cfg"] = cfg
+ ret["in_channels"] = in_channels
+ ret["mask_classification"] = mask_classification
+ 
+ ret["num_classes"] = cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES
+ ret["hidden_dim"] = cfg.MODEL.MASK_FORMER.HIDDEN_DIM
+ ret["num_queries"] = cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES
+ # Transformer parameters:
+ ret["nheads"] = cfg.MODEL.MASK_FORMER.NHEADS
+ ret["dim_feedforward"] = cfg.MODEL.MASK_FORMER.DIM_FEEDFORWARD
+
+ # NOTE: because we add learnable query features which requires supervision,
+ # we add minus 1 to decoder layers to be consistent with our loss
+ # implementation: that is, number of auxiliary losses is always
+ # equal to number of decoder layers. With learnable query features, the number of
+ # auxiliary losses equals number of decoders plus 1.
+ assert cfg.MODEL.MASK_FORMER.DEC_LAYERS >= 1
+ ret["dec_layers"] = cfg.MODEL.MASK_FORMER.DEC_LAYERS - 1
+ ret["pre_norm"] = cfg.MODEL.MASK_FORMER.PRE_NORM
+ ret["enforce_input_project"] = cfg.MODEL.MASK_FORMER.ENFORCE_INPUT_PROJ
+
+ ret["mask_dim"] = cfg.MODEL.SEM_SEG_HEAD.MASK_DIM
+
+ return ret
+
+ def forward(self, x, mask_features, mask = None):
+ # x is a list of multi-scale feature
+ assert len(x) == self.num_feature_levels
+
+ bt, c_m, h_m, w_m = mask_features.shape
+ bs = bt // (self.cfg.ENTITY.CROP_SAMPLE_NUM_TRAIN+1) if self.training else 1
+ # bs = bt // self.num_views if self.training else 1
+ t_m = bt // bs
+ mask_features_2d = mask_features
+ mask_features_3d = mask_features.view(bs, t_m, c_m, h_m, w_m)
+
+ src_2d, src_3d = [], []
+ pos_2d, pos_3d = [], []
+ size_list = []
+
+ # disable mask, it does not affect performance
+ del mask
+
+ # pdb.set_trace()
+ for i in range(self.num_feature_levels):
+ size_list.append(x[i].shape[-2:])
+ pos_2d_, pos_3d_ = self.pe_layer(x[i].view(bs, t_m, -1, size_list[-1][0], size_list[-1][1]))
+ 
+ pos_3d.append(pos_3d_.flatten(3))
+ src_3d.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
+
+ pos_2d.append(pos_2d_.flatten(2))
+ src_2d.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
+
+ # NTxCxHW => NxTxCxHW => (TxHW)xNxC
+ _, c, hw = src_3d[-1].shape
+ pos_3d[-1] = pos_3d[-1].view(bs, t_m, c, hw).permute(1, 3, 0, 2).flatten(0, 1)
+ src_3d[-1] = src_3d[-1].view(bs, t_m, c, hw).permute(1, 3, 0, 2).flatten(0, 1)
+
+ pos_2d[-1] = pos_2d[-1].permute(2,0,1)
+ src_2d[-1] = src_2d[-1].permute(2,0,1)
+
+ # QxNxC
+ query_embed_2d = self.query_embed.weight.unsqueeze(1).repeat(1, bt, 1)
+ output_2d = self.query_feat.weight.unsqueeze(1).repeat(1, bt, 1)
+
+ predictions_class_2d = []
+ predictions_mask_2d = []
+
+ # prediction heads on learnable query features
+ outputs_class_2d, outputs_mask_2d, attn_mask_2d, embedding_2d = self.forward_prediction_heads(output_2d, mask_features_2d, output_type="2d", attn_mask_target_size=size_list[0])
+ predictions_class_2d.append(outputs_class_2d)
+ predictions_mask_2d.append(outputs_mask_2d)
+ 
+ # pdb.set_trace()
+ for i in range(self.num_layers):
+ level_index = i % self.num_feature_levels
+ attn_mask_2d[torch.where(attn_mask_2d.sum(-1) == attn_mask_2d.shape[-1])] = False
+ # attention: cross-attention first
+ output_2d = self.transformer_cross_attention_layers[i](
+ output_2d, src_2d[level_index],
+ memory_mask=attn_mask_2d,
+ memory_key_padding_mask=None, # here we do not apply masking on padded region
+ pos=pos_2d[level_index], query_pos=query_embed_2d
+ )
+
+ output_2d = self.transformer_self_attention_layers[i](
+ output_2d, tgt_mask=None,
+ tgt_key_padding_mask=None,
+ query_pos=query_embed_2d
+ )
+ 
+ # FFN
+ output_2d = self.transformer_ffn_layers[i](
+ output_2d
+ )
+
+ outputs_class_2d, outputs_mask_2d, attn_mask_2d, embedding_2d = self.forward_prediction_heads(output_2d, mask_features_2d, output_type="2d", attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels])
+ predictions_class_2d.append(outputs_class_2d)
+ predictions_mask_2d.append(outputs_mask_2d)
+
+ assert len(predictions_class_2d) == self.num_layers + 1
+
+ out_2d = {
+ 'pred_logits': predictions_class_2d[-1],
+ 'pred_masks': predictions_mask_2d[-1],
+ 'aux_outputs': self._set_aux_loss(
+ predictions_class_2d if self.mask_classification else None, predictions_mask_2d
+ )
+ }
+
+ predictions_class_3d = []
+ predictions_mask_3d = []
+
+ query_embed_3d = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
+
+ output_3d = self.make_3d(output_2d, query_embed_2d, query_embed_3d)
+ 
+ # self.fused
+ outputs_class_3d, outputs_mask_3d, attn_mask_3d, embedding_3d = self.forward_prediction_heads(output_3d, mask_features_3d, output_type="3d", attn_mask_target_size=size_list[0])
+ predictions_class_3d.append(outputs_class_3d)
+ predictions_mask_3d.append(outputs_mask_3d)
+
+ for i in range(self.num_layers):
+ level_index = i % self.num_feature_levels
+ attn_mask_3d[torch.where(attn_mask_3d.sum(-1) == attn_mask_3d.shape[-1])] = False
+ ################# 3d (unified) #############
+ # attention: cross-attention first
+ output_3d = self.transformer_cross_attention_layers[i](
+ output_3d, src_3d[level_index],
+ memory_mask=attn_mask_3d,
+ memory_key_padding_mask=None, # here we do not apply masking on padded region
+ pos=pos_3d[level_index], query_pos=query_embed_3d
+ )
+
+ output_3d = self.transformer_self_attention_layers[i](
+ output_3d, tgt_mask=None,
+ tgt_key_padding_mask=None,
+ query_pos=query_embed_3d
+ )
+ 
+ output_3d = self.transformer_ffn_layers[i](
+ output_3d
+ )
+
+ outputs_class_3d, outputs_mask_3d, attn_mask_3d, embedding_3d = self.forward_prediction_heads(output_3d, mask_features_3d, output_type="3d", attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels])
+ predictions_class_3d.append(outputs_class_3d)
+ predictions_mask_3d.append(outputs_mask_3d)
+
+ # assert len(predictions_class_3d) == self.num_layers + 1
+
+ out_3d = {
+ 'pred_logits': predictions_class_3d[-1],
+ 'pred_masks': predictions_mask_3d[-1],
+ 'aux_outputs': self._set_aux_loss(
+ predictions_class_3d if self.mask_classification else None, predictions_mask_3d
+ ),
+ }
+
+ return out_2d, out_3d
+
+ def forward_prediction_heads(self, output, mask_features, output_type, attn_mask_target_size):
+ decoder_output = self.decoder_norm(output)
+ decoder_output = decoder_output.transpose(0, 1)
+ outputs_class = self.class_embed(decoder_output)
+ mask_embed = self.mask_embed(decoder_output)
+ if output_type == "3d":
+ outputs_mask = torch.einsum("bqc,btchw->bqthw", mask_embed, mask_features)
+ b, q, t, _, _ = outputs_mask.shape
+ # NOTE: prediction is of higher-resolution
+ # [B, Q, T, H, W] -> [B, Q, T*H*W] -> [B, h, Q, T*H*W] -> [B*h, Q, T*HW]
+ attn_mask = F.interpolate(outputs_mask.flatten(0, 1), size=attn_mask_target_size, mode="bilinear", align_corners=False).view(
+ b, q, t, attn_mask_target_size[0], attn_mask_target_size[1])
+ # must use bool type
+ # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged.
+ attn_mask = (attn_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5).bool()
+ attn_mask = attn_mask.detach()
+ elif output_type == "2d":
+ outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)
+ # NOTE: prediction is of higher-resolution
+ # [B, Q, H, W] -> [B, Q, H*W] -> [B, h, Q, H*W] -> [B*h, Q, HW]
+ attn_mask = F.interpolate(outputs_mask, size=attn_mask_target_size, mode="bilinear", align_corners=False)
+ # must use bool type
+ # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged.
+ attn_mask = (attn_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5).bool()
+ attn_mask = attn_mask.detach()
+ else:
+ raise "the output_type should be 2d or 3d"
+ 
+ return outputs_class, outputs_mask, attn_mask, decoder_output
+
+ @torch.jit.unused
+ def _set_aux_loss(self, outputs_class, outputs_seg_masks):
+ # this is a workaround to make torchscript happy, as torchscript
+ # doesn't support dictionary with non-homogeneous values, such
+ # as a dict having both a Tensor and a list.
+ if self.mask_classification:
+ return [
+ {"pred_logits": a, "pred_masks": b}
+ for a, b in zip(outputs_class[:-1], outputs_seg_masks[:-1])
+ ]
+ else:
+ return [{"pred_masks": b} for b in outputs_seg_masks[:-1]]

annotator/entityseg/mask2former/modeling/transformer_decoder/mask2former_transformer_decoder.py

@@ -0,0 +1,461 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/detr.py
+import logging
+import fvcore.nn.weight_init as weight_init
+from typing import Optional
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+from detectron2.config import configurable
+from detectron2.layers import Conv2d
+
+from .position_encoding import PositionEmbeddingSine
+from .maskformer_transformer_decoder import TRANSFORMER_DECODER_REGISTRY
+
+
+class SelfAttentionLayer(nn.Module):
+
+ def __init__(self, d_model, nhead, dropout=0.0,
+ activation="relu", normalize_before=False):
+ super().__init__()
+ self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+
+ self.norm = nn.LayerNorm(d_model)
+ self.dropout = nn.Dropout(dropout)
+
+ self.activation = _get_activation_fn(activation)
+ self.normalize_before = normalize_before
+
+ self._reset_parameters()
+ 
+ def _reset_parameters(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+
+ def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+ return tensor if pos is None else tensor + pos
+
+ def forward_post(self, tgt,
+ tgt_mask: Optional[Tensor] = None,
+ tgt_key_padding_mask: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ q = k = self.with_pos_embed(tgt, query_pos)
+ tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
+ key_padding_mask=tgt_key_padding_mask)[0]
+ tgt = tgt + self.dropout(tgt2)
+ tgt = self.norm(tgt)
+
+ return tgt
+
+ def forward_pre(self, tgt,
+ tgt_mask: Optional[Tensor] = None,
+ tgt_key_padding_mask: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ tgt2 = self.norm(tgt)
+ q = k = self.with_pos_embed(tgt2, query_pos)
+ tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+ key_padding_mask=tgt_key_padding_mask)[0]
+ tgt = tgt + self.dropout(tgt2)
+ 
+ return tgt
+
+ def forward(self, tgt,
+ tgt_mask: Optional[Tensor] = None,
+ tgt_key_padding_mask: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ if self.normalize_before:
+ return self.forward_pre(tgt, tgt_mask,
+ tgt_key_padding_mask, query_pos)
+ return self.forward_post(tgt, tgt_mask,
+ tgt_key_padding_mask, query_pos)
+
+
+class CrossAttentionLayer(nn.Module):
+
+ def __init__(self, d_model, nhead, dropout=0.0,
+ activation="relu", normalize_before=False):
+ super().__init__()
+ self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+
+ self.norm = nn.LayerNorm(d_model)
+ self.dropout = nn.Dropout(dropout)
+
+ self.activation = _get_activation_fn(activation)
+ self.normalize_before = normalize_before
+
+ self._reset_parameters()
+ 
+ def _reset_parameters(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+
+ def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+ return tensor if pos is None else tensor + pos
+
+ def forward_post(self, tgt, memory,
+ memory_mask: Optional[Tensor] = None,
+ memory_key_padding_mask: Optional[Tensor] = None,
+ pos: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
+ key=self.with_pos_embed(memory, pos),
+ value=memory, attn_mask=memory_mask,
+ key_padding_mask=memory_key_padding_mask)[0]
+ tgt = tgt + self.dropout(tgt2)
+ tgt = self.norm(tgt)
+ 
+ return tgt
+
+ def forward_pre(self, tgt, memory,
+ memory_mask: Optional[Tensor] = None,
+ memory_key_padding_mask: Optional[Tensor] = None,
+ pos: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ tgt2 = self.norm(tgt)
+ tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+ key=self.with_pos_embed(memory, pos),
+ value=memory, attn_mask=memory_mask,
+ key_padding_mask=memory_key_padding_mask)[0]
+ tgt = tgt + self.dropout(tgt2)
+
+ return tgt
+
+ def forward(self, tgt, memory,
+ memory_mask: Optional[Tensor] = None,
+ memory_key_padding_mask: Optional[Tensor] = None,
+ pos: Optional[Tensor] = None,
+ query_pos: Optional[Tensor] = None):
+ if self.normalize_before:
+ return self.forward_pre(tgt, memory, memory_mask,
+ memory_key_padding_mask, pos, query_pos)
+ return self.forward_post(tgt, memory, memory_mask,
+ memory_key_padding_mask, pos, query_pos)
+
+
+class FFNLayer(nn.Module):
+
+ def __init__(self, d_model, dim_feedforward=2048, dropout=0.0,
+ activation="relu", normalize_before=False):
+ super().__init__()
+ # Implementation of Feedforward model
+ self.linear1 = nn.Linear(d_model, dim_feedforward)
+ self.dropout = nn.Dropout(dropout)
+ self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+ self.norm = nn.LayerNorm(d_model)
+
+ self.activation = _get_activation_fn(activation)
+ self.normalize_before = normalize_before
+
+ self._reset_parameters()
+ 
+ def _reset_parameters(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+
+ def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+ return tensor if pos is None else tensor + pos
+
+ def forward_post(self, tgt):
+ tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+ tgt = tgt + self.dropout(tgt2)
+ tgt = self.norm(tgt)
+ return tgt
+
+ def forward_pre(self, tgt):
+ tgt2 = self.norm(tgt)
+ tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+ tgt = tgt + self.dropout(tgt2)
+ return tgt
+
+ def forward(self, tgt):
+ if self.normalize_before:
+ return self.forward_pre(tgt)
+ return self.forward_post(tgt)
+
+
+def _get_activation_fn(activation):
+ """Return an activation function given a string"""
+ if activation == "relu":
+ return F.relu
+ if activation == "gelu":
+ return F.gelu
+ if activation == "glu":
+ return F.glu
+ raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
+
+
+class MLP(nn.Module):
+ """ Very simple multi-layer perceptron (also called FFN)"""
+
+ def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+ super().__init__()
+ self.num_layers = num_layers
+ h = [hidden_dim] * (num_layers - 1)
+ self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+ def forward(self, x):
+ for i, layer in enumerate(self.layers):
+ x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+ return x
+
+
+@TRANSFORMER_DECODER_REGISTRY.register()
+class MultiScaleMaskedTransformerDecoder(nn.Module):
+
+ _version = 2
+
+ def _load_from_state_dict(
+ self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+ ):
+ version = local_metadata.get("version", None)
+ if version is None or version < 2:
+ # Do not warn if train from scratch
+ scratch = True
+ logger = logging.getLogger(__name__)
+ for k in list(state_dict.keys()):
+ newk = k
+ if "static_query" in k:
+ newk = k.replace("static_query", "query_feat")
+ if newk != k:
+ state_dict[newk] = state_dict[k]
+ del state_dict[k]
+ scratch = False
+
+ if not scratch:
+ logger.warning(
+ f"Weight format of {self.__class__.__name__} have changed! "
+ "Please upgrade your models. Applying automatic conversion now ..."
+ )
+
+ @configurable
+ def __init__(
+ self,
+ in_channels,
+ mask_classification=True,
+ *,
+ num_classes: int,
+ hidden_dim: int,
+ num_queries: int,
+ nheads: int,
+ dim_feedforward: int,
+ dec_layers: int,
+ pre_norm: bool,
+ mask_dim: int,
+ enforce_input_project: bool,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ in_channels: channels of the input features
+ mask_classification: whether to add mask classifier or not
+ num_classes: number of classes
+ hidden_dim: Transformer feature dimension
+ num_queries: number of queries
+ nheads: number of heads
+ dim_feedforward: feature dimension in feedforward network
+ enc_layers: number of Transformer encoder layers
+ dec_layers: number of Transformer decoder layers
+ pre_norm: whether to use pre-LayerNorm or not
+ mask_dim: mask feature dimension
+ enforce_input_project: add input project 1x1 conv even if input
+ channels and hidden dim is identical
+ """
+ super().__init__()
+
+ assert mask_classification, "Only support mask classification model"
+ self.mask_classification = mask_classification
+
+ # positional encoding
+ N_steps = hidden_dim // 2
+ self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
+ 
+ # define Transformer decoder here
+ self.num_heads = nheads
+ self.num_layers = dec_layers
+ self.transformer_self_attention_layers = nn.ModuleList()
+ self.transformer_cross_attention_layers = nn.ModuleList()
+ self.transformer_ffn_layers = nn.ModuleList()
+
+ for _ in range(self.num_layers):
+ self.transformer_self_attention_layers.append(
+ SelfAttentionLayer(
+ d_model=hidden_dim,
+ nhead=nheads,
+ dropout=0.0,
+ normalize_before=pre_norm,
+ )
+ )
+
+ self.transformer_cross_attention_layers.append(
+ CrossAttentionLayer(
+ d_model=hidden_dim,
+ nhead=nheads,
+ dropout=0.0,
+ normalize_before=pre_norm,
+ )
+ )
+
+ self.transformer_ffn_layers.append(
+ FFNLayer(
+ d_model=hidden_dim,
+ dim_feedforward=dim_feedforward,
+ dropout=0.0,
+ normalize_before=pre_norm,
+ )
+ )
+
+ self.decoder_norm = nn.LayerNorm(hidden_dim)
+
+ self.num_queries = num_queries
+ # learnable query features
+ self.query_feat = nn.Embedding(num_queries, hidden_dim)
+ # learnable query p.e.
+ self.query_embed = nn.Embedding(num_queries, hidden_dim)
+
+ # level embedding (we always use 3 scales)
+ self.num_feature_levels = 3
+ self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim)
+ self.input_proj = nn.ModuleList()
+ for _ in range(self.num_feature_levels):
+ if in_channels != hidden_dim or enforce_input_project:
+ self.input_proj.append(Conv2d(in_channels, hidden_dim, kernel_size=1))
+ weight_init.c2_xavier_fill(self.input_proj[-1])
+ else:
+ self.input_proj.append(nn.Sequential())
+
+ # output FFNs
+ if self.mask_classification:
+ self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
+ self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3)
+
+ @classmethod
+ def from_config(cls, cfg, in_channels, mask_classification):
+ ret = {}
+ ret["in_channels"] = in_channels
+ ret["mask_classification"] = mask_classification
+ 
+ ret["num_classes"] = cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES
+ ret["hidden_dim"] = cfg.MODEL.MASK_FORMER.HIDDEN_DIM
+ ret["num_queries"] = cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES
+ # Transformer parameters:
+ ret["nheads"] = cfg.MODEL.MASK_FORMER.NHEADS
+ ret["dim_feedforward"] = cfg.MODEL.MASK_FORMER.DIM_FEEDFORWARD
+
+ # NOTE: because we add learnable query features which requires supervision,
+ # we add minus 1 to decoder layers to be consistent with our loss
+ # implementation: that is, number of auxiliary losses is always
+ # equal to number of decoder layers. With learnable query features, the number of
+ # auxiliary losses equals number of decoders plus 1.
+ assert cfg.MODEL.MASK_FORMER.DEC_LAYERS >= 1
+ ret["dec_layers"] = cfg.MODEL.MASK_FORMER.DEC_LAYERS - 1
+ ret["pre_norm"] = cfg.MODEL.MASK_FORMER.PRE_NORM
+ ret["enforce_input_project"] = cfg.MODEL.MASK_FORMER.ENFORCE_INPUT_PROJ
+
+ ret["mask_dim"] = cfg.MODEL.SEM_SEG_HEAD.MASK_DIM
+
+ return ret
+
+ def forward(self, x, mask_features, mask = None):
+ # x is a list of multi-scale feature
+ assert len(x) == self.num_feature_levels
+ src = []
+ pos = []
+ size_list = []
+
+ # disable mask, it does not affect performance
+ del mask
+
+ for i in range(self.num_feature_levels):
+ size_list.append(x[i].shape[-2:])
+ pos.append(self.pe_layer(x[i], None).flatten(2))
+ src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
+
+ # flatten NxCxHxW to HWxNxC
+ pos[-1] = pos[-1].permute(2, 0, 1)
+ src[-1] = src[-1].permute(2, 0, 1)
+
+ _, bs, _ = src[0].shape
+
+ # QxNxC
+ query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
+ output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)
+
+ predictions_class = []
+ predictions_mask = []
+
+ # prediction heads on learnable query features
+ outputs_class, outputs_mask, attn_mask = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0])
+ predictions_class.append(outputs_class)
+ predictions_mask.append(outputs_mask)
+
+ for i in range(self.num_layers):
+ level_index = i % self.num_feature_levels
+ attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
+ # attention: cross-attention first
+ output = self.transformer_cross_attention_layers[i](
+ output, src[level_index],
+ memory_mask=attn_mask,
+ memory_key_padding_mask=None, # here we do not apply masking on padded region
+ pos=pos[level_index], query_pos=query_embed
+ )
+
+ output = self.transformer_self_attention_layers[i](
+ output, tgt_mask=None,
+ tgt_key_padding_mask=None,
+ query_pos=query_embed
+ )
+ 
+ # FFN
+ output = self.transformer_ffn_layers[i](
+ output
+ )
+
+ outputs_class, outputs_mask, attn_mask = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels])
+ predictions_class.append(outputs_class)
+ predictions_mask.append(outputs_mask)
+
+ assert len(predictions_class) == self.num_layers + 1
+
+ out = {
+ 'pred_logits': predictions_class[-1],
+ 'pred_masks': predictions_mask[-1],
+ 'aux_outputs': self._set_aux_loss(
+ predictions_class if self.mask_classification else None, predictions_mask
+ )
+ }
+ return out
+
+ def forward_prediction_heads(self, output, mask_features, attn_mask_target_size):
+ decoder_output = self.decoder_norm(output)
+ decoder_output = decoder_output.transpose(0, 1)
+ outputs_class = self.class_embed(decoder_output)
+ mask_embed = self.mask_embed(decoder_output)
+ outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)
+
+ # NOTE: prediction is of higher-resolution
+ # [B, Q, H, W] -> [B, Q, H*W] -> [B, h, Q, H*W] -> [B*h, Q, HW]
+ attn_mask = F.interpolate(outputs_mask, size=attn_mask_target_size, mode="bilinear", align_corners=False)
+ # must use bool type
+ # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged.
+ attn_mask = (attn_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5).bool()
+ attn_mask = attn_mask.detach()
+
+ return outputs_class, outputs_mask, attn_mask
+
+ @torch.jit.unused
+ def _set_aux_loss(self, outputs_class, outputs_seg_masks):
+ # this is a workaround to make torchscript happy, as torchscript
+ # doesn't support dictionary with non-homogeneous values, such
+ # as a dict having both a Tensor and a list.
+ if self.mask_classification:
+ return [
+ {"pred_logits": a, "pred_masks": b}
+ for a, b in zip(outputs_class[:-1], outputs_seg_masks[:-1])
+ ]
+ else:
+ return [{"pred_masks": b} for b in outputs_seg_masks[:-1]]

annotator/entityseg/mask2former/modeling/transformer_decoder/maskformer_transformer_decoder.py

@@ -0,0 +1,188 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/detr.py
+import fvcore.nn.weight_init as weight_init
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from detectron2.config import configurable
+from detectron2.layers import Conv2d
+from detectron2.utils.registry import Registry
+
+from .position_encoding import PositionEmbeddingSine
+from .transformer import Transformer
+
+
+TRANSFORMER_DECODER_REGISTRY = Registry("TRANSFORMER_MODULE")
+TRANSFORMER_DECODER_REGISTRY.__doc__ = """
+Registry for transformer module in MaskFormer.
+"""
+
+
+def build_transformer_decoder(cfg, in_channels, mask_classification=True):
+ """
+ Build a instance embedding branch from `cfg.MODEL.INS_EMBED_HEAD.NAME`.
+ """
+ name = cfg.MODEL.MASK_FORMER.TRANSFORMER_DECODER_NAME
+ return TRANSFORMER_DECODER_REGISTRY.get(name)(cfg, in_channels, mask_classification)
+
+
+@TRANSFORMER_DECODER_REGISTRY.register()
+class StandardTransformerDecoder(nn.Module):
+ @configurable
+ def __init__(
+ self,
+ in_channels,
+ mask_classification=True,
+ *,
+ num_classes: int,
+ hidden_dim: int,
+ num_queries: int,
+ nheads: int,
+ dropout: float,
+ dim_feedforward: int,
+ enc_layers: int,
+ dec_layers: int,
+ pre_norm: bool,
+ deep_supervision: bool,
+ mask_dim: int,
+ enforce_input_project: bool,
+ ):
+ """
+ NOTE: this interface is experimental.
+ Args:
+ in_channels: channels of the input features
+ mask_classification: whether to add mask classifier or not
+ num_classes: number of classes
+ hidden_dim: Transformer feature dimension
+ num_queries: number of queries
+ nheads: number of heads
+ dropout: dropout in Transformer
+ dim_feedforward: feature dimension in feedforward network
+ enc_layers: number of Transformer encoder layers
+ dec_layers: number of Transformer decoder layers
+ pre_norm: whether to use pre-LayerNorm or not
+ deep_supervision: whether to add supervision to every decoder layers
+ mask_dim: mask feature dimension
+ enforce_input_project: add input project 1x1 conv even if input
+ channels and hidden dim is identical
+ """
+ super().__init__()
+
+ self.mask_classification = mask_classification
+
+ # positional encoding
+ N_steps = hidden_dim // 2
+ self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
+
+ transformer = Transformer(
+ d_model=hidden_dim,
+ dropout=dropout,
+ nhead=nheads,
+ dim_feedforward=dim_feedforward,
+ num_encoder_layers=enc_layers,
+ num_decoder_layers=dec_layers,
+ normalize_before=pre_norm,
+ return_intermediate_dec=deep_supervision,
+ )
+
+ self.num_queries = num_queries
+ self.transformer = transformer
+ hidden_dim = transformer.d_model
+
+ self.query_embed = nn.Embedding(num_queries, hidden_dim)
+
+ if in_channels != hidden_dim or enforce_input_project:
+ self.input_proj = Conv2d(in_channels, hidden_dim, kernel_size=1)
+ weight_init.c2_xavier_fill(self.input_proj)
+ else:
+ self.input_proj = nn.Sequential()
+ self.aux_loss = deep_supervision
+
+ # output FFNs
+ if self.mask_classification:
+ self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
+ self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3)
+
+ @classmethod
+ def from_config(cls, cfg, in_channels, mask_classification):
+ ret = {}
+ ret["in_channels"] = in_channels
+ ret["mask_classification"] = mask_classification
+
+ ret["num_classes"] = cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES
+ ret["hidden_dim"] = cfg.MODEL.MASK_FORMER.HIDDEN_DIM
+ ret["num_queries"] = cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES
+ # Transformer parameters:
+ ret["nheads"] = cfg.MODEL.MASK_FORMER.NHEADS
+ ret["dropout"] = cfg.MODEL.MASK_FORMER.DROPOUT
+ ret["dim_feedforward"] = cfg.MODEL.MASK_FORMER.DIM_FEEDFORWARD
+ ret["enc_layers"] = cfg.MODEL.MASK_FORMER.ENC_LAYERS
+ ret["dec_layers"] = cfg.MODEL.MASK_FORMER.DEC_LAYERS
+ ret["pre_norm"] = cfg.MODEL.MASK_FORMER.PRE_NORM
+ ret["deep_supervision"] = cfg.MODEL.MASK_FORMER.DEEP_SUPERVISION
+ ret["enforce_input_project"] = cfg.MODEL.MASK_FORMER.ENFORCE_INPUT_PROJ
+
+ ret["mask_dim"] = cfg.MODEL.SEM_SEG_HEAD.MASK_DIM
+
+ return ret
+
+ def forward(self, x, mask_features, mask=None):
+ if mask is not None:
+ mask = F.interpolate(mask[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
+ pos = self.pe_layer(x, mask)
+
+ src = x
+ hs, memory = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos)
+
+ if self.mask_classification:
+ outputs_class = self.class_embed(hs)
+ out = {"pred_logits": outputs_class[-1]}
+ else:
+ out = {}
+
+ if self.aux_loss:
+ # [l, bs, queries, embed]
+ mask_embed = self.mask_embed(hs)
+ outputs_seg_masks = torch.einsum("lbqc,bchw->lbqhw", mask_embed, mask_features)
+ out["pred_masks"] = outputs_seg_masks[-1]
+ out["aux_outputs"] = self._set_aux_loss(
+ outputs_class if self.mask_classification else None, outputs_seg_masks
+ )
+ else:
+ # FIXME h_boxes takes the last one computed, keep this in mind
+ # [bs, queries, embed]
+ mask_embed = self.mask_embed(hs[-1])
+ outputs_seg_masks = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)
+ out["pred_masks"] = outputs_seg_masks
+ return out
+
+ @torch.jit.unused
+ def _set_aux_loss(self, outputs_class, outputs_seg_masks):
+ # this is a workaround to make torchscript happy, as torchscript
+ # doesn't support dictionary with non-homogeneous values, such
+ # as a dict having both a Tensor and a list.
+ if self.mask_classification:
+ return [
+ {"pred_logits": a, "pred_masks": b}
+ for a, b in zip(outputs_class[:-1], outputs_seg_masks[:-1])
+ ]
+ else:
+ return [{"pred_masks": b} for b in outputs_seg_masks[:-1]]
+
+
+class MLP(nn.Module):
+ """Very simple multi-layer perceptron (also called FFN)"""
+
+ def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+ super().__init__()
+ self.num_layers = num_layers
+ h = [hidden_dim] * (num_layers - 1)
+ self.layers = nn.ModuleList(
+ nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+ )
+
+ def forward(self, x):
+ for i, layer in enumerate(self.layers):
+ x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+ return x

annotator/entityseg/mask2former/modeling/transformer_decoder/position_encoding.py

@@ -0,0 +1,134 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# # Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/position_encoding.py
+"""
+Various positional encodings for the transformer.
+"""
+import math
+
+import torch
+from torch import nn
+
+
+class PositionEmbeddingSine(nn.Module):
+ """
+ This is a more standard version of the position embedding, very similar to the one
+ used by the Attention is all you need paper, generalized to work on images.
+ """
+
+ def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+ super().__init__()
+ self.num_pos_feats = num_pos_feats
+ self.temperature = temperature
+ self.normalize = normalize
+ if scale is not None and normalize is False:
+ raise ValueError("normalize should be True if scale is passed")
+ if scale is None:
+ scale = 2 * math.pi
+ self.scale = scale
+
+ def forward(self, x, mask=None):
+ if mask is None:
+ mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
+ not_mask = ~mask
+ y_embed = not_mask.cumsum(1, dtype=torch.float32)
+ x_embed = not_mask.cumsum(2, dtype=torch.float32)
+ if self.normalize:
+ eps = 1e-6
+ y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+ x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+ dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+ # dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+ dim_t = self.temperature ** (2 * (torch.div(dim_t, 2, rounding_mode="trunc")) / self.num_pos_feats)
+
+ pos_x = x_embed[:, :, :, None] / dim_t
+ pos_y = y_embed[:, :, :, None] / dim_t
+ pos_x = torch.stack(
+ (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4
+ ).flatten(3)
+ pos_y = torch.stack(
+ (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4
+ ).flatten(3)
+ pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+ return pos
+ 
+ def __repr__(self, _repr_indent=4):
+ head = "Positional encoding " + self.__class__.__name__
+ body = [
+ "num_pos_feats: {}".format(self.num_pos_feats),
+ "temperature: {}".format(self.temperature),
+ "normalize: {}".format(self.normalize),
+ "scale: {}".format(self.scale),
+ ]
+ # _repr_indent = 4
+ lines = [head] + [" " * _repr_indent + line for line in body]
+ return "\n".join(lines)
+
+class PositionEmbeddingSine3D2D(nn.Module):
+ """
+ This is a more standard version of the position embedding, very similar to the one
+ used by the Attention is all you need paper, generalized to work on images.
+ """
+
+ def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+ super().__init__()
+ self.num_pos_feats = num_pos_feats
+ self.temperature = temperature
+ self.normalize = normalize
+ if scale is not None and normalize is False:
+ raise ValueError("normalize should be True if scale is passed")
+ if scale is None:
+ scale = 2 * math.pi
+ self.scale = scale
+
+ def forward(self, x, mask=None):
+ ## b, t, c, h, w
+ assert x.dim()==5, f"{x.shape} should be a 5-dimensional Tensor, got {x.dim()}-dimensional Tensor instead"
+ if mask is None:
+ mask = torch.zeros((x.size(0), x.size(1), x.size(3), x.size(4)), device=x.device, dtype=torch.bool)
+ not_mask = ~mask
+ z_embed = not_mask.cumsum(1, dtype=torch.float32)
+ y_embed = not_mask.cumsum(2, dtype=torch.float32)
+ x_embed = not_mask.cumsum(3, dtype=torch.float32)
+ if self.normalize:
+ eps = 1e-6
+ z_embed = z_embed / (z_embed[:, -1:, :, :] + eps) * self.scale 
+ y_embed = y_embed / (y_embed[:, :, -1:, :] + eps) * self.scale
+ x_embed = x_embed / (x_embed[:, :, :, -1:] + eps) * self.scale
+
+ dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+ # dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+ dim_t = self.temperature ** (2 * (torch.div(dim_t, 2, rounding_mode="trunc")) / self.num_pos_feats)
+
+ dim_t_z = torch.arange((self.num_pos_feats * 2), dtype=torch.float32, device=x.device)
+ # dim_t_z = self.temperature ** (2 * (dim_t_z // 2) / (self.num_pos_feats * 2))
+ dim_t_z = self.temperature ** (2 * (torch.div(dim_t_z, 2, rounding_mode="trunc")) / (self.num_pos_feats*2))
+
+ pos_x = x_embed[:, :, :, :, None] / dim_t
+ pos_y = y_embed[:, :, :, :, None] / dim_t
+ pos_z = z_embed[:, :, :, :, None] / dim_t_z
+
+ pos_x = torch.stack(
+ (pos_x[:, :, :, :, 0::2].sin(), pos_x[:, :, :, :, 1::2].cos()), dim=5
+ ).flatten(4)
+ pos_y = torch.stack(
+ (pos_y[:, :, :, :, 0::2].sin(), pos_y[:, :, :, :, 1::2].cos()), dim=5
+ ).flatten(4)
+ pos_z = torch.stack(
+ (pos_z[:, :, :, :, 0::2].sin(), pos_z[:, :, :, :, 1::2].cos()), dim=5
+ ).flatten(4)
+ pos2d = torch.cat((pos_y, pos_x), dim=4).permute(0, 1, 4, 2, 3).flatten(0,1)
+ pos3d = (torch.cat((pos_y, pos_x), dim=4) + pos_z).permute(0, 1, 4, 2, 3)
+ return pos2d, pos3d
+ 
+ def __repr__(self, _repr_indent=4):
+ head = "Positional encoding " + self.__class__.__name__
+ body = [
+ "num_pos_feats: {}".format(self.num_pos_feats),
+ "temperature: {}".format(self.temperature),
+ "normalize: {}".format(self.normalize),
+ "scale: {}".format(self.scale),
+ ]
+ # _repr_indent = 4
+ lines = [head] + [" " * _repr_indent + line for line in body]
+ return "\n".join(lines)