Unverified 提交 0d8a1842 authored 作者: ELHoussineT's avatar ELHoussineT 提交者: GitHub

Add `crops = results.crops()` dictionary (#4676)

* adding get cropped functionality * Add target logic in existing functions * Crops cleanup * Add dictionary keys: conf, cls, box * Bug fixes - avoid return after first image Co-authored-by: 's avatarGlenn Jocher <glenn.jocher@ultralytics.com>
上级 8e94bf62
...@@ -365,6 +365,7 @@ class Detections: ...@@ -365,6 +365,7 @@ class Detections:
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, save_dir=Path('')):
crops = []
for i, (im, pred) in enumerate(zip(self.imgs, self.pred)): for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '
if pred.shape[0]: if pred.shape[0]:
...@@ -376,7 +377,9 @@ class Detections: ...@@ -376,7 +377,9 @@ class Detections:
for *box, conf, cls in reversed(pred): # xyxy, confidence, class for *box, conf, cls in reversed(pred): # xyxy, confidence, class
label = f'{self.names[int(cls)]} {conf:.2f}' label = f'{self.names[int(cls)]} {conf:.2f}'
if crop: if crop:
save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i]) file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
'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, color=colors(cls))
im = annotator.im im = annotator.im
...@@ -395,6 +398,10 @@ class Detections: ...@@ -395,6 +398,10 @@ class Detections:
LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}") LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
if render: if render:
self.imgs[i] = np.asarray(im) self.imgs[i] = np.asarray(im)
if crop:
if save:
LOGGER.info(f'Saved results to {save_dir}\n')
return crops
def print(self): def print(self):
self.display(pprint=True) # print results self.display(pprint=True) # print results
...@@ -408,10 +415,9 @@ class Detections: ...@@ -408,10 +415,9 @@ class Detections:
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, save_dir=save_dir) # save results
def crop(self, 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) # increment save_dir save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
self.display(crop=True, save_dir=save_dir) # crop results return self.display(crop=True, save=save, save_dir=save_dir) # crop results
LOGGER.info(f'Saved results to {save_dir}\n')
def render(self): def render(self):
self.display(render=True) # render results self.display(render=True) # render results
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论