Spaces:
Runtime error
Runtime error
File size: 4,016 Bytes
0c7479d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
# Copyright (c) Facebook, Inc. and its affiliates.
import logging
from typing import Callable, Dict, List, Optional, Tuple, Union
import math
import copy
import numpy as np
import torch
import torch.distributed as dist
from torch import nn
from torch.nn import functional as F
from fvcore.nn import sigmoid_focal_loss_jit, giou_loss, smooth_l1_loss
import fvcore.nn.weight_init as weight_init
from detectron2.config import configurable
from detectron2.data.detection_utils import get_fed_loss_cls_weights
from detectron2.layers import ShapeSpec, batched_nms, cat, cross_entropy, nonzero_tuple
from detectron2.modeling.box_regression import Box2BoxTransform, _dense_box_regression_loss
from detectron2.structures import Boxes, Instances, BitMasks, pairwise_iou, pairwise_ioa
from detectron2.utils.events import get_event_storage
from detectron2.modeling.roi_heads.fast_rcnn import FastRCNNOutputLayers
class TexeEmbedClassifier(nn.Module):
def __init__(
self,
input_shape: ShapeSpec,
zs_weight_dim: int = 1024,
norm_weight: bool = True,
norm_temperature: float = 50.0,
):
super().__init__()
if isinstance(input_shape, int): # some backward compatibility
input_shape = ShapeSpec(channels=input_shape)
input_size = input_shape.channels * (input_shape.width or 1) * (input_shape.height or 1)
self.norm_weight = norm_weight
self.norm_temperature = norm_temperature
self.linear = nn.Linear(input_size, zs_weight_dim)
def forward(self, x, text_embed):
x = self.linear(x)
if self.norm_weight:
x = self.norm_temperature * F.normalize(x, p=2, dim=1)
x = torch.mm(x, text_embed)
return x
class VLMFastRCNNOutputLayers(nn.Module):
def __init__(
self,
input_shape: ShapeSpec,
box2box_transform,
use_sigmoid_ce: bool = True,
test_score_thresh: float = 0.0,
test_nms_thresh: float = 0.5,
test_topk_per_image: int = 100,
):
super().__init__()
if isinstance(input_shape, int): # some backward compatibility
input_shape = ShapeSpec(channels=input_shape)
self.box2box_transform = box2box_transform
self.use_sigmoid_ce = use_sigmoid_ce
self.test_score_thresh = test_score_thresh
self.test_nms_thresh = test_nms_thresh
self.test_topk_per_image = test_topk_per_image
input_size = input_shape.channels * \
(input_shape.width or 1) * (input_shape.height or 1)
# bbox_pred
self.bbox_pred = nn.Sequential(
nn.Linear(input_size, input_size),
nn.ReLU(inplace=True),
nn.Linear(input_size, 4)
)
# cls_score
self.cls_score = TexeEmbedClassifier(input_shape)
def forward(self, x, text_embed):
if x.dim() > 2:
x = torch.flatten(x, start_dim=1)
cls_scores = self.cls_score(x, text_embed)
proposal_deltas = self.bbox_pred(x)
return cls_scores, proposal_deltas
def predict_boxes(self, predictions, proposals):
if not len(proposals):
return []
_, proposal_deltas = predictions
num_prop_per_image = [len(p) for p in proposals]
proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0)
predict_boxes = self.box2box_transform.apply_deltas(
proposal_deltas,
proposal_boxes,
) # Nx(KxB)
return predict_boxes.split(num_prop_per_image)
def predict_probs(self, predictions, proposals):
cls_scores, _ = predictions
num_inst_per_image = [len(p) for p in proposals]
cls_scores = cls_scores.split(num_inst_per_image, dim=0)
final_scores = []
for cls_score in cls_scores:
final_score = cls_score.sigmoid() if self.use_sigmoid_ce else F.softmax(cls_score, dim=-1)
final_scores.append(final_score)
return final_scores
|