♻️ [Refactor] Vec2Box, Align with YOLOv7 structure
Browse files
yolo/utils/bounding_box_utils.py
CHANGED
|
@@ -265,12 +265,12 @@ class BoxMatcher:
|
|
| 265 |
|
| 266 |
|
| 267 |
class Vec2Box:
|
| 268 |
-
def __init__(self, model: YOLO, image_size, device):
|
| 269 |
self.device = device
|
| 270 |
|
| 271 |
-
if hasattr(
|
| 272 |
-
logger.info(f"🈶 Found stride of model {
|
| 273 |
-
self.strides =
|
| 274 |
else:
|
| 275 |
logger.info("🧸 Found no stride of model, performed a dummy test for auto-anchor size")
|
| 276 |
self.strides = self.create_auto_anchor(model, image_size)
|
|
|
|
| 265 |
|
| 266 |
|
| 267 |
class Vec2Box:
|
| 268 |
+
def __init__(self, model: YOLO, anchor_cfg: AnchorConfig, image_size, device):
|
| 269 |
self.device = device
|
| 270 |
|
| 271 |
+
if hasattr(anchor_cfg, "strides"):
|
| 272 |
+
logger.info(f"🈶 Found stride of model {anchor_cfg.strides}")
|
| 273 |
+
self.strides = anchor_cfg.strides
|
| 274 |
else:
|
| 275 |
logger.info("🧸 Found no stride of model, performed a dummy test for auto-anchor size")
|
| 276 |
self.strides = self.create_auto_anchor(model, image_size)
|