yolov8 class working

micromind-toolkit · Jan 11, 2024 · 5d3fbd3 · 5d3fbd3
1 parent 2ba9318
commit 5d3fbd3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/micromind/networks/yolo.py b/micromind/networks/yolo.py
@@ -554,7 +554,9 @@ class YOLOv8(nn.Module):
     def __init__(self, w, r, d, num_classes=80):
         super().__init__()
         self.net = Darknet(w, r, d)
-        self.fpn = Yolov8Neck(w, r, d)
+        self.fpn = Yolov8Neck(
+            filters=[int(256 * w), int(512 * w), int(512 * w * r)], d=d
+        )
         self.head = DetectionHead(
             num_classes, filters=(int(256 * w), int(512 * w), int(512 * w * r))
         )