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
Example #2
0
 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
Example #3
0
    def _init_test_anchor_generator(self):
        strides = [4]
        anchor_grids = [[6, 14]]
        anchor_generator = AnchorGenerator(strides, anchor_grids)

        return anchor_generator
Example #4
0
 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