zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

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: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
ELHoussineT 2021-09-08 19:42:28 +02:00 committed by GitHub
parent 8e94bf62d9
commit 0d8a184237
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
models/common.py
View File

@ -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:
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:
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:
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