Add cv2 support

detect.py

@@ -182,7 +182,7 @@ def wrap_frozen_graph(gd, inputs, outputs):
             s += '%gx%g ' % img.shape[2:]  # print string
             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
             imc = im0.copy() if save_crop else im0  # for save_crop
-            annotator = Annotator(im0, line_width=line_thickness)
+            annotator = Annotator(im0, line_width=line_thickness, pil=False)
             if len(det):
                 # Rescale boxes from img_size to im0 size
                 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
@@ -267,7 +267,7 @@ def parse_opt():
     parser.add_argument('--project', default='runs/detect', help='save results to project/name')
     parser.add_argument('--name', default='exp', help='save results to project/name')
     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
-    parser.add_argument('--line-thickness', default=4, type=int, help='bounding box thickness (pixels)')
+    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
     parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
     parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
     parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')

utils/general.py

@@ -122,9 +122,10 @@ def is_pip():
     return 'site-packages' in Path(__file__).absolute().parts
 
 
-def is_ascii(str=''):
+def is_ascii(s=''):
     # Is string composed of all ASCII (no UTF) characters?
-    return len(str.encode().decode('ascii', 'ignore')) == len(str)
+    s = str(s)  # convert to str() in case of None, etc.
+    return len(s.encode().decode('ascii', 'ignore')) == len(s)
 
 
 def emojis(str=''):

utils/plots.py

@@ -96,35 +96,51 @@ def plot_one_box(box, im, color=(128, 128, 128), txt_color=(255, 255, 255), labe
 
 class Annotator:
     # YOLOv5 PIL Annotator class
-    def __init__(self, im, line_width=4, font_size=None, font='Arial.ttf'):
+    def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=True):
         assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
-        self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
-        self.draw = ImageDraw.Draw(self.im)
-        s = sum(self.im.size) / 2  # mean shape
-        f = font_size or max(round(s * 0.035), 12)
-        self.lw = line_width or max(round(s * 0.005), 2)  # line width
-        try:
-            self.font = ImageFont.truetype(font, size=f)
-        except:  # download TTF
-            url = "https://github.com/ultralytics/yolov5/releases/download/v1.0/" + font
-            torch.hub.download_url_to_file(url, font)
-            self.font = ImageFont.truetype(font, size=f)
-        self.fh = self.font.getsize('a')[1] - 3  # font height
+        self.pil = pil
+        if self.pil:  # use PIL
+            self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
+            self.draw = ImageDraw.Draw(self.im)
+            s = sum(self.im.size) / 2  # mean shape
+            f = font_size or max(round(s * 0.035), 12)
+            try:
+                self.font = ImageFont.truetype(font, size=f)
+            except:  # download TTF
+                url = "https://github.com/ultralytics/yolov5/releases/download/v1.0/" + font
+                torch.hub.download_url_to_file(url, font)
+                self.font = ImageFont.truetype(font, size=f)
+            self.fh = self.font.getsize('a')[1] - 3  # font height
+        else:  # use cv2
+            self.im = im
+        s = sum(self.im.shape) / 2  # mean shape
+        self.lw = line_width or max(round(s * 0.003), 2)  # line width
 
     def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
         # Add one xyxy box to image with label
-        self.draw.rectangle(box, width=self.lw, outline=color)  # box
-        if label:
-            w = self.font.getsize(label)[0]  # text width
-            self.draw.rectangle([box[0], box[1] - self.fh, box[0] + w + 1, box[1] + 1], fill=color)
-            self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls')
+        if self.pil or not is_ascii(label):
+            self.draw.rectangle(box, width=self.lw, outline=color)  # box
+            if label:
+                w = self.font.getsize(label)[0]  # text width
+                self.draw.rectangle([box[0], box[1] - self.fh, box[0] + w + 1, box[1] + 1], fill=color)
+                self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls')
+        else:  # cv2
+            c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
+            cv2.rectangle(self.im, c1, c2, color, thickness=self.lw, lineType=cv2.LINE_AA)
+            if label:
+                tf = max(self.lw - 1, 1)  # font thickness
+                w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0]
+                c2 = c1[0] + w, c1[1] - h - 3
+                cv2.rectangle(self.im, c1, c2, color, -1, cv2.LINE_AA)  # filled
+                cv2.putText(self.im, label, (c1[0], c1[1] - 2), 0, self.lw / 3, txt_color, thickness=tf,
+                            lineType=cv2.LINE_AA)
 
     def rectangle(self, xy, fill=None, outline=None, width=1):
-        # Add rectangle to image
+        # Add rectangle to image (PIL-only)
         self.draw.rectangle(xy, fill, outline, width)
 
     def text(self, xy, text, txt_color=(255, 255, 255)):
-        # Add text to image
+        # Add text to image (PIL-only)
         w, h = self.font.getsize(text)  # text width, height
         self.draw.text((xy[0], xy[1] - h + 1), text, fill=txt_color, font=self.font)