Python draws a rectangular frame in a picture

Functional description

1. Draw three types of rectangular frames in a picture, which are solid lines, dotted lines and dashed lines;
2. Realize the continuous drawing of multiple pictures.

Instructions for use

1. "bbox_data.xlsx" is the coordinates of the rectangular box in the picture, corresponding to [xmin, ymin, xmax, ymax]
   
   as shown in the figure above, the leftmost is the number, followed by the coordinates of three rectangular boxes
2. Use digital format for pictures, such as "1.jpg" and "2.jpg", so that they can correspond to each other according to the numbers in "bbox_data.xlsx".

the code

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