python 在图片中绘制矩形框

功能描述

1. 在一张图片中绘制三种类型的矩形框，分别为实线、点线和虚线；
2. 实现连续地多张图片图片的绘制。

使用说明

1. “bbox_data.xlsx”为矩形框在图片中的坐标，对应[xmin, ymin, xmax, ymax]
   
   如上图所示，最左边是编号，后面跟了三个矩形框的坐标
2. 图片使用数字格式，如"1.jpg"、“2.jpg”，这样才能根据“bbox_data.xlsx”中的编号相互对应。

代码

import os

import cv2
import numpy as np
import pandas as pd


def drawline(img, pt1, pt2, color, thickness=1, style='dotted', gap=10):
    dist = ((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2) ** .5
    pts = []
    for i in np.arange(0, dist, gap):
        r = i / dist
        x = int((pt1[0] * (1 - r) + pt2[0] * r) + .5)
        y = int((pt1[1] * (1 - r) + pt2[1] * r) + .5)
        p = (x, y)
        pts.append(p)

    if style == 'dotted':
        for p in pts:
            cv2.circle(img, p, thickness, color, -1)
    else:
        s = pts[0]
        e = pts[0]
        i = 0
        for p in pts:
            s = e
            e = p
            if i % 2 == 1:
                cv2.line(img, s, e, color, thickness)
            i += 1


def drawRect(img, pt1, pt2, color, thickness, style):
    drawline(img, (pt1[0], pt1[1]), (pt2[0], pt1[1]), color, thickness, style)
    drawline(img, (pt2[0], pt1[1]), (pt2[0], pt2[1]), color, thickness, style)
    drawline(img, (pt2[0], pt2[1]), (pt1[0], pt2[1]), color, thickness, style)
    drawline(img, (pt1[0], pt2[1]), (pt1[0], pt1[1]), color, thickness, style)


bbox_data = pd.read_excel(r'./bbox_data.xlsx', 'Sheet1')

img_source = './row_img/'
img_data = os.listdir(img_source)
img_data.sort()
count = 1
for idx in range(0, bbox_data.shape[0], 3):
    for j in range(idx, idx + 3):
        img = cv2.imread(os.path.join(img_source, img_data[j]))
        # 读取bbox的第j行
        tmp_data = bbox_data.iloc[j]
        bbox1_kalman = list(tmp_data[1: 5])
        if bbox1_kalman[0] != 0:
            img = cv2.rectangle(img, (bbox1_kalman[0], bbox1_kalman[1]), (bbox1_kalman[2], bbox1_kalman[3]), (0, 0, 0),6)

        bbox2_kcf = list(tmp_data[5: 9])
        if bbox2_kcf[0] != 0:
            drawRect(img, (bbox2_kcf[0], bbox2_kcf[1]), (bbox2_kcf[2], bbox2_kcf[3]), (0, 0, 0), thickness=6,
                     style='dotted')
        bbox3_tld = list(tmp_data[9: 13])
        if bbox3_tld[0] != 0:
            drawRect(img, (bbox3_tld[0], bbox3_tld[1]), (bbox3_tld[2],  bbox3_tld[3]), (0, 0, 0), thickness=6, style='dotted2')
        cv2.imwrite(r'./drawRecResult/' + str(count) + '.jpg', img)
        count += 1