def _init_test_backbone_with_fpn(self):
backbone_name = 'darknet_s_r3_1'
depth_multiple = 0.33
width_multiple = 0.5
backbone_with_fpn = darknet_pan_backbone(backbone_name, depth_multiple,
width_multiple)
return backbone_with_fpn
def __init__(self, checkpoint_path: str, version: str = "r6.0"):
super().__init__()
model_info = load_from_ultralytics(checkpoint_path, version=version)
backbone_name = f"darknet_{model_info['size']}_{version.replace('.', '_')}"
depth_multiple = model_info["depth_multiple"]
width_multiple = model_info["width_multiple"]
use_p6 = model_info["use_p6"]
backbone = darknet_pan_backbone(
backbone_name,
depth_multiple,
width_multiple,
version=version,
use_p6=use_p6,
)
num_classes = model_info["num_classes"]
anchor_generator = AnchorGenerator(model_info["strides"],
model_info["anchor_grids"])
post_process = LogitsDecoder(model_info["strides"])
model = YOLO(
backbone,
num_classes,
anchor_generator=anchor_generator,
post_process=post_process,
)
model.load_state_dict(model_info["state_dict"])
self.model = model
self.num_classes = num_classes