glenn-jocher commited on
Commit
a0a4adf
·
unverified ·
1 Parent(s): bc3ed95

Add PyTorch Hub `results.save(labels=False)` option (#7129)

Browse files

Resolves https://github.com/ultralytics/yolov5/issues/388#issuecomment-1077121821

Files changed (1) hide show
  1. models/common.py +9 -9
models/common.py CHANGED
@@ -131,7 +131,7 @@ class C3(nn.Module):
131
  c_ = int(c2 * e) # hidden channels
132
  self.cv1 = Conv(c1, c_, 1, 1)
133
  self.cv2 = Conv(c1, c_, 1, 1)
134
- self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
135
  self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
136
  # self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))
137
 
@@ -589,7 +589,7 @@ class Detections:
589
  self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
590
  self.s = shape # inference BCHW shape
591
 
592
- def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
593
  crops = []
594
  for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
595
  s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string
@@ -606,7 +606,7 @@ class Detections:
606
  crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
607
  'im': save_one_box(box, im, file=file, save=save)})
608
  else: # all others
609
- annotator.box_label(box, label, color=colors(cls))
610
  im = annotator.im
611
  else:
612
  s += '(no detections)'
@@ -633,19 +633,19 @@ class Detections:
633
  LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' %
634
  self.t)
635
 
636
- def show(self):
637
- self.display(show=True) # show results
638
 
639
- def save(self, save_dir='runs/detect/exp'):
640
  save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
641
- self.display(save=True, save_dir=save_dir) # save results
642
 
643
  def crop(self, save=True, save_dir='runs/detect/exp'):
644
  save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
645
  return self.display(crop=True, save=save, save_dir=save_dir) # crop results
646
 
647
- def render(self):
648
- self.display(render=True) # render results
649
  return self.imgs
650
 
651
  def pandas(self):
 
131
  c_ = int(c2 * e) # hidden channels
132
  self.cv1 = Conv(c1, c_, 1, 1)
133
  self.cv2 = Conv(c1, c_, 1, 1)
134
+ self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)
135
  self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
136
  # self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))
137
 
 
589
  self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
590
  self.s = shape # inference BCHW shape
591
 
592
+ def display(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path('')):
593
  crops = []
594
  for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
595
  s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string
 
606
  crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
607
  'im': save_one_box(box, im, file=file, save=save)})
608
  else: # all others
609
+ annotator.box_label(box, label if labels else '', color=colors(cls))
610
  im = annotator.im
611
  else:
612
  s += '(no detections)'
 
633
  LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' %
634
  self.t)
635
 
636
+ def show(self, labels=True):
637
+ self.display(show=True, labels=labels) # show results
638
 
639
+ def save(self, labels=True, save_dir='runs/detect/exp'):
640
  save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
641
+ self.display(save=True, labels=labels, save_dir=save_dir) # save results
642
 
643
  def crop(self, save=True, save_dir='runs/detect/exp'):
644
  save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
645
  return self.display(crop=True, save=save, save_dir=save_dir) # crop results
646
 
647
+ def render(self, labels=True):
648
+ self.display(render=True, labels=labels) # render results
649
  return self.imgs
650
 
651
  def pandas(self):