【python脚本】在图像中标注目标，并保存结果

说明
【功能】在图像中用鼠标拖一个矩形，并将这个矩形框的像素坐标保存到文件
【输入参数】

• image_path: 图像的根路径
• output_file: 输出文件路径，如 /xx/xx/res.txt

import argparse
import cv2
import os

class RectangularSelection:
    def __init__(self, image_path, output_file):
        self.image_path = image_path
        self.output_file = output_file

    def run(self):
        # 1. check and prepare
        if not os.path.exists(self.image_path):
            print('\33[31mError: invalid image path "{}"\33[0m'.format(self.image_path))
            return
        ofs = open(self.output_file, 'w')

        # 2. traverse and save result to file
        names = self.obtain_file_names(self.image_path, 'jpg')
        # print(names)
        for timestamp, name in names:
            exit = self.click_rectangle(timestamp, name)
            if exit == True:
                break

            # check and save result
            if len(self.rec_pts) > 0 and len(self.rec_pts) % 2 == 0:
                while True:
                    check = self.check_rectangle(timestamp, name)
                    if check == True:
                        ofs.write(str(timestamp))
                        for pt in self.rec_pts:
                            ofs.write(" " + str(pt[0]) + " " + str(pt[1]))
                        ofs.write("\n")
                        break
                    else:
                        exit = self.click_rectangle(timestamp, name)
                        if exit == True:
                            exit(0)

    def obtain_file_names(self, image_path, suffix):
        res = []
        for _, _, names in os.walk(image_path):
            names.sort(key=lambda n: int(n.split('.')[0]))
            for name in names:
                if name.split('.')[-1] == suffix:
                    res.append(((float(name.split('.')[0]) / 1e6), os.path.join(image_path, name)))
            break
        return res

    def click_and_crop(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            self.rec_pts.append((x, y))
        elif event == cv2.EVENT_LBUTTONUP:
            self.rec_pts.append((x, y))
            # draw a rectangle around the region of interest
            if len(self.rec_pts) >= 2:
                cv2.rectangle(self.image, self.rec_pts[-2], self.rec_pts[-1], (0, 255, 0), 2)
                cv2.imshow("image", self.image)

    def click_rectangle(self, timestamp, name):
        exit = False
        self.image = cv2.imread(name)
        clone = self.image.copy()
        cv2.namedWindow("image")
        cv2.setMouseCallback("image", self.click_and_crop)
        self.rec_pts = []
        while True:
            cv2.putText(self.image, str(timestamp), (5,50), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
            cv2.putText(self.image, "please click left top and right bottom (c: ok and next, q: exit)", (5,100),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
            cv2.imshow("image", self.image)
            key = cv2.waitKey(1) & 0xFF
            if key == ord("r"):
                self.image = clone.copy()
                self.rec_pts = []
            elif key == ord("c"):
                break
            elif key == ord("q"):
                exit = True
                break
        cv2.destroyAllWindows()
        return exit

    def check_rectangle(self, timestamp, name):
        check = False
        cv2.namedWindow("check")
        img = cv2.imread(name)
        for i in range(0, len(self.rec_pts), 2):
            cv2.rectangle(img, self.rec_pts[i], self.rec_pts[i + 1], (0, 0, 255), 2)
        while True:
            cv2.putText(img, str(timestamp), (5,50), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
            cv2.putText(img, "please check (c: ok and next, x: do it again)", (5,100),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
            cv2.imshow("check", img)
            key = cv2.waitKey(1) & 0xFF
            if key == ord("c"):
                check = True
                break
            elif key == ord("x"):
                check = False
                break
        return check

if __name__=="__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--image_path", required=True, help="Path to the image")
    ap.add_argument("-o", "--output_file", required=True, help="output_path")
    args = ap.parse_args()

    rec_select = RectangularSelection(args.image_path, args.output_file)
    rec_select.run()