File size: 3,512 Bytes
561a469 |
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 |
import numpy as np
from transformers.models.deformable_detr.modeling_deformable_detr import DeformableDetrMLPPredictionHead
import torch.nn as nn
import torch
def PairDetr(model, num_queries, num_classes):
in_features = model.class_embed[0].in_features
model.model.query_position_embeddings = nn.Embedding(num_queries, 512)
class_embed = nn.Linear(in_features, num_classes)
bbox_embed = DeformableDetrMLPPredictionHead(
input_dim=256, hidden_dim=256, output_dim=8, num_layers=3
)
model.class_embed = nn.ModuleList([class_embed for _ in range(6)])
model.bbox_embed = nn.ModuleList([bbox_embed for _ in range(6)])
return model
def inverse_sigmoid(x, eps=1e-5):
x = x.clamp(min=0, max=1)
x1 = x.clamp(min=eps)
x2 = (1 - x).clamp(min=eps)
return torch.log(x1 / x2)
def forward(model,
pixel_values,
pixel_mask=None,
decoder_attention_mask=None,
encoder_outputs=None,
inputs_embeds=None,
decoder_inputs_embeds=None,
labels=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,) -> torch.Tensor:
return_dict = return_dict if return_dict is not None else model.config.use_return_dict
outputs = model.model(
pixel_values,
pixel_mask=pixel_mask,
decoder_attention_mask=decoder_attention_mask,
encoder_outputs=encoder_outputs,
inputs_embeds=inputs_embeds,
decoder_inputs_embeds=decoder_inputs_embeds,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
hidden_states = outputs.intermediate_hidden_states if return_dict else outputs[2]
init_reference = outputs.init_reference_points if return_dict else outputs[0]
inter_references = outputs.intermediate_reference_points if return_dict else outputs[3]
outputs_classes = []
outputs_coords = []
cons = inverse_sigmoid(init_reference)
for level in range(hidden_states.shape[1]):
if level == 0:
reference = init_reference
else:
reference = inter_references[:, level - 1]
reference = inverse_sigmoid(reference)
outputs_class = model.class_embed[level](hidden_states[:, level])
delta_bbox = model.bbox_embed[level](hidden_states[:, level])
if reference.shape[-1] == 4:
delta_bbox[..., :4] += reference
outputs_coord_logits = delta_bbox
elif reference.shape[-1] == 2:
delta_bbox[..., :2] += reference
delta_bbox[..., 4:6] += cons
outputs_coord_logits = delta_bbox
else:
raise ValueError(f"reference.shape[-1] should be 4 or 2, but got {reference.shape[-1]}")
outputs_coord = outputs_coord_logits.sigmoid()
outputs_classes.append(outputs_class)
outputs_coords.append(outputs_coord)
outputs_class = torch.stack(outputs_classes, dim=1)
outputs_coord = torch.stack(outputs_coords, dim=1)
logits = outputs_class[:, -1]
pred_boxes = outputs_coord[:, -1]
dict_outputs = {
"logits":logits,
"pred_boxes": pred_boxes,
"init_reference_points": outputs.init_reference_points,
}
return dict_outputs
|