see-- / keras-centernet

A Keras implementation of CenterNet with pre-trained model (unofficial)
MIT License
337 stars 84 forks source link

draw_box() is too slow #13

Closed tak-s closed 4 years ago

tak-s commented 5 years ago

To draw detected object labels, convert to Pillow image object. But this process is too heavy cost. I changed draw_box() in utils.py like below,

  def draw_box(self, img, x1, y1, x2, y2, cl):
    cl = int(cl)
    x1, y1, x2, y2 = int(round(float(x1))), int(round(float(y1))), int(round(float(x2))), int(round(float(y2)))
    h = img.shape[0]
    width = max(1, int(h * 0.006))
    name = self.coco_names[cl].split()[-1]
    bgr_color = get_rgb_color(cl, len(self.coco_names))[::-1]
    # bounding box
    #cv2.rectangle(img, (x1, y1), (x2, y2), bgr_color, width)
    cv2.rectangle(img, (x1, y1), (x2, y2), bgr_color, 2)
    # font background
    #font_width = len(name) * self.char_width
    #cv2.rectangle(img, (x1 - math.ceil(width / 2), y1 - self.font_size), (x1 + font_width, y1), bgr_color, -1)
    # text
    #pil_img = Image.fromarray(img[..., ::-1])
    #draw = ImageDraw.Draw(pil_img)
    #draw.text((x1 + width, y1 - self.font_size), name, font=self.font, fill=(0, 0, 0, 255))
    #img = np.array(pil_img)[..., ::-1].copy()
    # draw text shadow
    cv2.putText(img, name, (x1+width+1, y1-6+1), cv2.FONT_HERSHEY_PLAIN, 1.2, (0,0,0), 2, cv2.LINE_AA)
    # draw text shadow
    cv2.putText(img, name, (x1+width, y1-6), cv2.FONT_HERSHEY_PLAIN, 1.2, bgr_color, 2, cv2.LINE_AA)

    return img
see-- commented 5 years ago

Thanks for sharing your faster version. I agree that the Pillow conversions are a bit wasteful. How does this image look like with your changes?

tak-s commented 5 years ago

Hi. Original code is

fn=assets/demo.jpg, TIME=0.4375
fn=assets/hp_demo.jpg, TIME=0.2267
fn=assets/demo2.jpg, TIME=0.4976

My code is

fn=assets/demo.jpg, TIME=0.1760
fn=assets/hp_demo.jpg, TIME=0.1513
fn=assets/demo2.jpg, TIME=0.1756

'TIME' is the time from cv2.imread to cv2.imwrite(). My GPU is GeForce RTX 2070.