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
def _init_test_anchor_generator(self):
strides = self._get_strides()
anchor_grids = self._get_anchor_grids()
anchor_generator = AnchorGenerator(strides, anchor_grids)
return anchor_generator
def _init_test_anchor_generator(self):
strides = [4]
anchor_grids = [[6, 14]]
anchor_generator = AnchorGenerator(strides, anchor_grids)
return anchor_generator
def _init_test_anchor_generator(self, use_p6=False):
strides = self._get_strides(use_p6)
anchor_grids = self._get_anchor_grids(use_p6)
anchor_generator = AnchorGenerator(strides, anchor_grids)
return anchor_generator