Commit
•
0d8a184
1
Parent(s):
8e94bf6
Add `crops = results.crops()` dictionary (#4676)
Browse files* 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>
- models/common.py +11 -5
models/common.py
CHANGED
@@ -365,6 +365,7 @@ class Detections:
|
|
365 |
self.s = shape # inference BCHW shape
|
366 |
|
367 |
def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
|
|
|
368 |
for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
|
369 |
str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '
|
370 |
if pred.shape[0]:
|
@@ -376,7 +377,9 @@ class Detections:
|
|
376 |
for *box, conf, cls in reversed(pred): # xyxy, confidence, class
|
377 |
label = f'{self.names[int(cls)]} {conf:.2f}'
|
378 |
if crop:
|
379 |
-
|
|
|
|
|
380 |
else: # all others
|
381 |
annotator.box_label(box, label, color=colors(cls))
|
382 |
im = annotator.im
|
@@ -395,6 +398,10 @@ class Detections:
|
|
395 |
LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
|
396 |
if render:
|
397 |
self.imgs[i] = np.asarray(im)
|
|
|
|
|
|
|
|
|
398 |
|
399 |
def print(self):
|
400 |
self.display(pprint=True) # print results
|
@@ -408,10 +415,9 @@ class Detections:
|
|
408 |
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
|
409 |
self.display(save=True, save_dir=save_dir) # save results
|
410 |
|
411 |
-
def crop(self, save_dir='runs/detect/exp'):
|
412 |
-
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True)
|
413 |
-
self.display(crop=True, save_dir=save_dir) # crop results
|
414 |
-
LOGGER.info(f'Saved results to {save_dir}\n')
|
415 |
|
416 |
def render(self):
|
417 |
self.display(render=True) # render results
|
|
|
365 |
self.s = shape # inference BCHW shape
|
366 |
|
367 |
def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
|
368 |
+
crops = []
|
369 |
for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
|
370 |
str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '
|
371 |
if pred.shape[0]:
|
|
|
377 |
for *box, conf, cls in reversed(pred): # xyxy, confidence, class
|
378 |
label = f'{self.names[int(cls)]} {conf:.2f}'
|
379 |
if crop:
|
380 |
+
file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
|
381 |
+
crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
|
382 |
+
'im': save_one_box(box, im, file=file, save=save)})
|
383 |
else: # all others
|
384 |
annotator.box_label(box, label, color=colors(cls))
|
385 |
im = annotator.im
|
|
|
398 |
LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
|
399 |
if render:
|
400 |
self.imgs[i] = np.asarray(im)
|
401 |
+
if crop:
|
402 |
+
if save:
|
403 |
+
LOGGER.info(f'Saved results to {save_dir}\n')
|
404 |
+
return crops
|
405 |
|
406 |
def print(self):
|
407 |
self.display(pprint=True) # print results
|
|
|
415 |
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
|
416 |
self.display(save=True, save_dir=save_dir) # save results
|
417 |
|
418 |
+
def crop(self, save=True, save_dir='runs/detect/exp'):
|
419 |
+
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
|
420 |
+
return self.display(crop=True, save=save, save_dir=save_dir) # crop results
|
|
|
421 |
|
422 |
def render(self):
|
423 |
self.display(render=True) # render results
|