glenn-jocher commited on
Commit
d929bb6
1 Parent(s): efa7a91

Implement default class names (#1609)

Browse files
Files changed (1) hide show
  1. 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
- sys.path.append('./') # to run '$ python *.py' files in subdirectories
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