glenn-jocher
commited on
Commit
•
d929bb6
1
Parent(s):
efa7a91
Implement default class names (#1609)
Browse files- models/yolo.py +5 -4
models/yolo.py
CHANGED
@@ -1,16 +1,16 @@
|
|
1 |
import argparse
|
2 |
import logging
|
3 |
-
import math
|
4 |
import sys
|
5 |
from copy import deepcopy
|
6 |
from pathlib import Path
|
7 |
|
8 |
-
|
9 |
-
logger = logging.getLogger(__name__)
|
10 |
-
|
11 |
import torch
|
12 |
import torch.nn as nn
|
13 |
|
|
|
|
|
|
|
14 |
from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, NMS, autoShape
|
15 |
from models.experimental import MixConv2d, CrossConv, C3
|
16 |
from utils.autoanchor import check_anchor_order
|
@@ -82,6 +82,7 @@ class Model(nn.Module):
|
|
82 |
logger.info('Overriding model.yaml nc=%g with nc=%g' % (self.yaml['nc'], nc))
|
83 |
self.yaml['nc'] = nc # override yaml value
|
84 |
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out
|
|
|
85 |
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
86 |
|
87 |
# Build strides, anchors
|
|
|
1 |
import argparse
|
2 |
import logging
|
|
|
3 |
import sys
|
4 |
from copy import deepcopy
|
5 |
from pathlib import Path
|
6 |
|
7 |
+
import math
|
|
|
|
|
8 |
import torch
|
9 |
import torch.nn as nn
|
10 |
|
11 |
+
sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
12 |
+
logger = logging.getLogger(__name__)
|
13 |
+
|
14 |
from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, NMS, autoShape
|
15 |
from models.experimental import MixConv2d, CrossConv, C3
|
16 |
from utils.autoanchor import check_anchor_order
|
|
|
82 |
logger.info('Overriding model.yaml nc=%g with nc=%g' % (self.yaml['nc'], nc))
|
83 |
self.yaml['nc'] = nc # override yaml value
|
84 |
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out
|
85 |
+
self.names = [str(i) for i in range(self.yaml['nc'])] # default names
|
86 |
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
87 |
|
88 |
# Build strides, anchors
|