27 lines
631 B
Python
27 lines
631 B
Python
import torch
|
|
|
|
from ..modules.text_grounding_net import UniFusion
|
|
|
|
|
|
def get_positionnet_default_params():
|
|
return {
|
|
"in_dim": 768,
|
|
"mid_dim": 3072,
|
|
"out_dim": 768,
|
|
"test_drop_boxes": False,
|
|
"test_drop_masks": True,
|
|
"test_drop_points": False,
|
|
"test_drop_scribbles": True,
|
|
"train_add_boxes": True,
|
|
"train_add_masks": True,
|
|
"train_add_points": True,
|
|
"train_add_scribbles": True,
|
|
"use_seperate_tokenizer": True,
|
|
}
|
|
|
|
|
|
def prepare_positionnet(checkpoint, params) -> torch.nn.Module:
|
|
model = UniFusion(**params)
|
|
model.load_state_dict(checkpoint)
|
|
return model
|