Browse Source

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

Resolves https://github.com/ultralytics/yolov5/issues/388#issuecomment-1077121821
modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
a0a4adf6de
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 9 deletions
  1. +9
    -9
      models/common.py

+ 9
- 9
models/common.py View File

c_ = int(c2 * e) # hidden channels c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1) self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1) self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n))) self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
# self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n))) # self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))


self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms) self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
self.s = shape # inference BCHW shape self.s = shape # inference BCHW shape


def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
def display(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path('')):
crops = [] crops = []
for i, (im, pred) in enumerate(zip(self.imgs, self.pred)): for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string
crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label, crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
'im': save_one_box(box, im, file=file, save=save)}) 'im': save_one_box(box, im, file=file, save=save)})
else: # all others else: # all others
annotator.box_label(box, label, color=colors(cls))
annotator.box_label(box, label if labels else '', color=colors(cls))
im = annotator.im im = annotator.im
else: else:
s += '(no detections)' s += '(no detections)'
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' %
self.t) self.t)


def show(self):
self.display(show=True) # show results
def show(self, labels=True):
self.display(show=True, labels=labels) # show results


def save(self, save_dir='runs/detect/exp'):
def save(self, labels=True, save_dir='runs/detect/exp'):
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
self.display(save=True, save_dir=save_dir) # save results
self.display(save=True, labels=labels, save_dir=save_dir) # save results


def crop(self, save=True, save_dir='runs/detect/exp'): def crop(self, save=True, save_dir='runs/detect/exp'):
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
return self.display(crop=True, save=save, save_dir=save_dir) # crop results return self.display(crop=True, save=save, save_dir=save_dir) # crop results


def render(self):
self.display(render=True) # render results
def render(self, labels=True):
self.display(render=True, labels=labels) # render results
return self.imgs return self.imgs


def pandas(self): def pandas(self):

Loading…
Cancel
Save