def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
config = omegaconf.OmegaConf.merge(
effdet.get_efficientdet_config(model_name=MODEL_NAME),
omegaconf.OmegaConf.create({
"min_level":
MIN_LEVEL,
"max_level":
MAX_LEVEL,
"num_levels":
NUM_LEVELS,
"max_detection_points":
MAX_DETECTION_POINTS,
}),
)
self.model = effdet.create_model_from_config(
config=config,
num_classes=NUM_CLASSES,
bench_task="",
pretrained=False,
pretrained_backbone=False,
image_size=(INPUT_HEIGHT, INPUT_WIDTH),
max_det_per_image=MAX_DET_PER_IMAGE,
)
self.config = {
"anchors": effdet.anchors.Anchors.from_config(self.model.config)
}
self.resize = torchvision.transforms.Resize(size=(INPUT_HEIGHT,
INPUT_WIDTH))
def __init__(
self,
annotation_config=mds.COCOAnnotationConfiguration90,
model_name: str = "tf_efficientdet_d0",
pretrained_backbone: bool = True,
pretrained_top: bool = False,
device="cpu",
resize_method: detector.ResizeMethod = "fit",
**kwargs,
):
super().__init__(device=device, resize_method=resize_method)
config = effdet.get_efficientdet_config(model_name=model_name)
if kwargs:
config = omegaconf.OmegaConf.merge( # type: ignore
config,
omegaconf.OmegaConf.create(kwargs),
)
self.annotation_config = annotation_config
self.model = effdet.create_model_from_config(
config=config,
num_classes=len(annotation_config),
bench_task="",
pretrained=pretrained_top,
pretrained_backbone=pretrained_backbone,
).to(self.device)
self.model_name = model_name
self.set_input_shape(width=config.image_size[1],
height=config.image_size[0])
def get_effdet_model(cfg, pretrained=True, task='train'):
base_config = get_efficientdet_config(cfg.train.backbone_det)
base_config.image_size = (cfg.data.img_size, cfg.data.img_size)
model = create_model_from_config(base_config,
bench_task=task,
bench_labeler=True,
num_classes=cfg.data.num_classes,
pretrained=pretrained)
return model
def model(
backbone: EfficientDetBackboneConfig,
num_classes: int,
img_size: int,
**kwargs,
) -> nn.Module:
"""Creates the efficientdet model specified by `model_name`.
The model implementation is by Ross Wightman, original repo
[here](https://github.com/rwightman/efficientdet-pytorch).
# Arguments
backbone: Specifies the backbone to use create the model. For pretrained models, check
[this](https://github.com/rwightman/efficientdet-pytorch#models) table.
num_classes: Number of classes of your dataset (including background).
img_size: Image size that will be fed to the model. Must be squared and
divisible by 128.
# Returns
A PyTorch model.
"""
model_name = backbone.model_name
config = get_efficientdet_config(model_name=model_name)
config.image_size = (img_size,
img_size) if isinstance(img_size, int) else img_size
model_bench = create_model_from_config(
config,
bench_task="train",
bench_labeler=True,
num_classes=num_classes - 1,
pretrained=backbone.pretrained,
**kwargs,
)
# TODO: Break down param groups for backbone
def param_groups_fn(model: nn.Module) -> List[List[nn.Parameter]]:
unwrapped = unwrap_bench(model)
layers = [
unwrapped.backbone,
unwrapped.fpn,
nn.Sequential(unwrapped.class_net, unwrapped.box_net),
]
param_groups = [list(layer.parameters()) for layer in layers]
check_all_model_params_in_groups2(model, param_groups)
return param_groups
model_bench.param_groups = MethodType(param_groups_fn, model_bench)
return model_bench
def get_model(
self,
model_name="tf_efficientdet_d0",
impactonly=False,
seqmode=False,
fullsizeimage=False,
anchor_scale=4,
):
model_name = model_name
config = get_efficientdet_config(model_name)
if fullsizeimage:
# (H, W) size % 128 = 0 on each dim
config.image_size = (640, 1280)
# config.image_size = (1280, 1280)
else:
config.image_size = (512, 512)
config.anchor_scale = anchor_scale
# config.norm_kwargs = dict(eps=0.001, momentum=0.01)
if impactonly:
num_classes = 1
else:
num_classes = 2
model = create_model_from_config(
config=config,
bench_task="train",
num_classes=num_classes,
pretrained=True,
checkpoint_path="",
checkpoint_ema=False,
bench_labeler=True,
)
if seqmode == True:
model.model.backbone.conv_stem = timm.models.layers.Conv2dSame(
9, 32, kernel_size=(3, 3), stride=(2, 2), bias=False)
return model