PyTorch Hub results.render() (#1897)

models/common.py

@@ -1,6 +1,7 @@
 # This file contains modules common to various models
 
 import math
+
 import numpy as np
 import requests
 import torch
@@ -240,7 +241,7 @@ class Detections:
         self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
         self.n = len(self.pred)
 
-    def display(self, pprint=False, show=False, save=False):
+    def display(self, pprint=False, show=False, save=False, render=False):
         colors = color_list()
         for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
             str = f'Image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
@@ -248,19 +249,21 @@ class Detections:
                 for c in pred[:, -1].unique():
                     n = (pred[:, -1] == c).sum()  # detections per class
                     str += f'{n} {self.names[int(c)]}s, '  # add to string
-                if show or save:
+                if show or save or render:
                     img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img  # from np
                     for *box, conf, cls in pred:  # xyxy, confidence, class
                         # str += '%s %.2f, ' % (names[int(cls)], conf)  # label
                         ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10])  # plot
+            if pprint:
+                print(str)
+            if show:
+                img.show(f'Image {i}')  # show
             if save:
                 f = f'results{i}.jpg'
                 str += f"saved to '{f}'"
                 img.save(f)  # save
-            if show:
-                img.show(f'Image {i}')  # show
-            if pprint:
-                print(str)
+            if render:
+                self.imgs[i] = np.asarray(img)
 
     def print(self):
         self.display(pprint=True)  # print results
@@ -271,6 +274,10 @@ class Detections:
     def save(self):
         self.display(save=True)  # save results
 
+    def render(self):
+        self.display(render=True)  # render results
+        return self.imgs
+
     def __len__(self):
         return self.n