97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
import torch
|
|
|
|
|
|
class GroundingNetInput:
|
|
def __init__(self):
|
|
self.set = False
|
|
self.return_att_masks = False
|
|
self.image_size = 64
|
|
self.return_att_masks32 = False
|
|
|
|
def prepare(self, batch, image_size=64, device=None, dtype=None, return_att_masks=False):
|
|
"""
|
|
batch should be the output from dataset.
|
|
Please define here how to process the batch and prepare the
|
|
input only for the ground tokenizer.
|
|
"""
|
|
output = {}
|
|
self.set = True
|
|
self.return_att_masks = return_att_masks
|
|
|
|
boxes = batch['boxes']
|
|
masks = batch['masks']
|
|
positive_embeddings = batch["prompts"]
|
|
|
|
if self.return_att_masks:
|
|
assert 'att_masks' in batch
|
|
att_masks = batch['att_masks']
|
|
|
|
scribbles = batch['scribbles']
|
|
polygons = batch['polygons']
|
|
self.dim_scribbles = scribbles.shape[-1]
|
|
self.dim_polygons = polygons.shape[-1]
|
|
# NOTE: New Seg
|
|
segs = batch["segments"]
|
|
self.dim_segs = segs.shape[-1]
|
|
points = batch["points"]
|
|
|
|
self.batch, self.max_box, self.in_dim = positive_embeddings.shape
|
|
self.device = positive_embeddings.device
|
|
self.dtype = positive_embeddings.dtype
|
|
|
|
output = {
|
|
"boxes": boxes,
|
|
"masks": masks,
|
|
"prompts": positive_embeddings,
|
|
}
|
|
output["scribbles"] = scribbles
|
|
output["polygons"] = polygons
|
|
output["segments"] = segs
|
|
output["points"] = points
|
|
|
|
if self.return_att_masks:
|
|
output['att_masks'] = att_masks
|
|
return output
|
|
|
|
def get_null_input(self, batch=None, latent_width=64, latent_height=64, device=None, dtype=None):
|
|
"""
|
|
Guidance for training (drop) or inference,
|
|
please define the null input for the grounding tokenizer
|
|
"""
|
|
|
|
assert self.set, "not set yet, cannot call this funcion"
|
|
batch = self.batch if batch is None else batch
|
|
device = self.device if device is None else device
|
|
dtype = self.dtype if dtype is None else dtype
|
|
|
|
boxes = torch.zeros(batch, self.max_box, 4,).type(dtype).to(device)
|
|
masks = torch.zeros(batch, self.max_box).type(dtype).to(device)
|
|
# NOTE: New Seg
|
|
segs = torch.zeros(batch, self.max_box, self.dim_segs,
|
|
self.dim_segs).type(dtype).to(device)
|
|
|
|
scribbles = torch.zeros(batch, self.max_box,
|
|
self.dim_scribbles).type(dtype).to(device)
|
|
polygons = torch.zeros(batch, self.max_box,
|
|
self.dim_polygons).type(dtype).to(device)
|
|
points = torch.zeros(batch, self.max_box, 2).type(dtype).to(device)
|
|
|
|
positive_embeddings = torch.zeros(
|
|
batch, self.max_box, self.in_dim).type(dtype).to(device)
|
|
|
|
output = {
|
|
"boxes": boxes,
|
|
"masks": masks,
|
|
"prompts": positive_embeddings,
|
|
}
|
|
output["scribbles"] = scribbles
|
|
output["polygons"] = polygons
|
|
output["segments"] = segs
|
|
output["points"] = points
|
|
|
|
if self.return_att_masks:
|
|
att_masks = torch.zeros(batch, self.max_box, latent_width, latent_height).type(
|
|
dtype).to(device) # TODO Order width/height
|
|
output['att_masks'] = att_masks
|
|
return output
|