henry000 commited on
Commit
fd35390
·
1 Parent(s): 7a913f7

♻️ [Refactor] Vec2Box, Align with YOLOv7 structure

Browse files
Files changed (1) hide show
  1. yolo/utils/bounding_box_utils.py +4 -4
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(model, "strides") and getattr(model, "strides"):
272
- logger.info(f"🈶 Found stride of model {model.strides}")
273
- self.strides = model.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)