python tkinter视频坐标标注工具

1、输入视频地址、点击加载按钮、选择模式

2、手动标注后点击右下角对应按钮，保存至当前目录下label.txt文件内

import cv2
from tkinter import *
from tkinter import messagebox
from PIL import Image, ImageTk

class LabelTool():
    def __init__(self, master):
        self.points = []
        self.parent = master
        self.parent.title("label")
        self.frame = Frame(self.parent)
        self.frame.pack(fill=BOTH, expand=1)
        self.parent.resizable(width=FALSE, height=FALSE)
        self.STATE = {'click': 1, 'x': 0, 'y': 0}

        self.label = Label(self.frame, text="请输入视频流地址：", font=('黑体', 10))
        self.label.grid(row=0, column=0, sticky=W)
        self.entry = Entry(self.frame)
        self.entry.grid(row=0, column=0)

        self.mainPanel = Canvas(self.frame, cursor='tcross')
        self.mainPanel.bind("<Button-1>", self.mouseClick)
        self.mainPanel.bind("<Motion>", self.mouseMove)
        self.mainPanel.grid(row=1, column=0, rowspan=8, sticky=N)

        self.ctrPanel = Frame(self.frame)
        self.ctrPanel.grid(row=9, column=0, columnspan=1, sticky=W + E)
        self.disp = Label(self.ctrPanel, text='', fg='red', font=('Arial', 12))
        self.disp.pack(side=LEFT)
        self.returnBtn = Button(self.ctrPanel, text="返回", command=self.stepBack, font=('黑体', 15))
        self.returnBtn.pack(side=RIGHT)

        self.mode = StringVar(self.parent)
        self.mode.set('1')
        affineMode = Radiobutton(self.frame, text='Affine Mode', value='1', variable=self.mode, command=self.clearPoints, font=('Arial', 15))
        lineMode = Radiobutton(self.frame, text='Line Mode', value='2', variable=self.mode, command=self.clearPoints, font=('Arial', 15))
        roadMode = Radiobutton(self.frame, text='Road Mode', value='3', variable=self.mode, command=self.clearPoints, font=('Arial', 15))
        roiMode = Radiobutton(self.frame, text='Roi Mode', value='4', variable=self.mode, command=self.clearPoints, font=('Arial', 15))

        affineMode.grid(row=0, column=1, sticky=W + N)
        lineMode.grid(row=0, column=2, sticky=W + N)
        roadMode.grid(row=1, column=1, sticky=W + N)
        roiMode.grid(row=1, column=2, sticky=W + N)

        self.loadBtn = Button(self.frame, text="Load Video", command=self.canva, font=('Arial', 12), bg='#00FA9A')
        self.loadBtn.grid(row=2, column=1, sticky=W + N + E)
        self.saveAffineBtn = Button(self.frame, text="Save to Affine", command=self.saveFile, font=('Arial', 12), bg="#C0C0C0")
        self.saveAffineBtn.grid(row=2, column=2, sticky=W + N + E)
        self.saveLine1Btn = Button(self.frame, text="Save to Line1", command=self.saveFile, font=('Arial', 12), bg="#FFD700")
        self.saveLine1Btn.grid(row=3, column=1, sticky=W + N + E)
        self.saveLine2Btn = Button(self.frame, text="Save to Line2", command=self.saveFile, font=('Arial', 12), bg="#FFFSEE")
        self.saveLine2Btn.grid(row=3, column=2, sticky=W + N + E)
        self.saveRoad1Btn = Button(self.frame, text="Save to Lane1", command=self.saveFile, font=('Arial', 12), bg="#B0E0E6")
        self.saveRoad1Btn.grid(row=4, column=1, sticky=W + N + E)
        self.saveRoad2Btn = Button(self.frame, text="Save to Lane2", command=self.saveFile, font=('Arial', 12), bg='#F5DEB3')
        self.saveRoad2Btn.grid(row=4, column=2, sticky=W + N + E)
        self.saveRoad3Btn = Button(self.frame, text="Save to Lane3", command=self.saveFile, font=('Arial', 12), bg='#7FFF00')
        self.saveRoad3Btn.grid(row=5, column=1, sticky=W + N + E)
        self.RoiBtn = Button(self.frame, text="Save to ROI", command=self.saveFile, font=('Arial', 12), bg='#FF00FF')
        self.RoiBtn.grid(row=5, column=2, sticky=W + N + E)

        self.text = Text(self.frame, height=8, width=40, fg='#FF0000', font=('Arial', 15), bg='#DEB887')
        self.text.grid(row=6, column=1, columnspan=2)

    def canva(self):
        frames = 0
        self.flag = 1
        self.points = []
        url = self.entry.get()
        if url == '0' or url == '':
            url = 0
        cap = cv2.VideoCapture(url)
        # assert cap.isOpened(), 'Cannot capture source'
        assert cap.isOpened(), self.warn()
        while True:
            frames += 1
            ret, frame = cap.read()
            if ret and frames%3 == 0:
                frames = 0
                img = cv2.resize(frame, (576, 324))
                if len(self.points) == 1:
                    cv2.circle(img, self.points[0], 6, (0, 255, 200), thickness=-1)

                if self.mode.get() in ['1', '3', '4']:
                    for i in range(len(self.points)-1):
                        cv2.circle(img, self.points[i], 6, (0, 255, 200), thickness=-1)
                        cv2.circle(img, self.points[i+1], 6, (50, 255, 200), thickness=-1)
                        if (i + 2) % 4 == 0:
                            cv2.line(img, self.points[i-2], self.points[i+1], (0, 0, 255), thickness=2)
                        if (i + 2) % 4 != 1 and (i + 1) % 4 != 0:
                            cv2.line(img, self.points[i], self.points[i+1], (0, 0, 255), thickness=2)
                    self.flag = (not(len(self.points) % 4 == 0))
                    if len(self.points) and self.flag == 1:
                        cv2.line(img, (self.STATE['x'],  self.STATE['y']), self.points[-1], (0, 255, 0), thickness=2)
                        cv2.circle(img, (self.STATE['x'],  self.STATE['y']), 5, (0, 255, 0), thickness=-1)

                if self.mode.get() == '2':
                    for i in range(len(self.points)-1):
                        cv2.circle(img, self.points[i], 6, (0, 255, 200), thickness=-1)
                        cv2.circle(img, self.points[i+1], 6, (50, 255, 200), thickness=-1)
                        if (i + 2) % 2 != 1 and (i + 1) % 2 != 0:
                            cv2.line(img, self.points[i], self.points[i+1], (0, 0, 255), thickness=2)
                    self.flag = (not(len(self.points) % 2 == 0))
                    if len(self.points) and self.flag == 1:
                        cv2.line(img, (self.STATE['x'],  self.STATE['y']), self.points[-1], (0, 255, 0), thickness=2)
                        cv2.circle(img, (self.STATE['x'],  self.STATE['y']), 5, (0, 255, 0), thickness=-1)

                img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGBA))
                tkImage = ImageTk.PhotoImage(image=img)
                self.mainPanel.config(width=tkImage.width(), height=tkImage.height())
                self.mainPanel.create_image(0, 0, anchor='nw', image=tkImage)
                self.parent.update()
                self.parent.after(100)
        cap.release()

    def mouseMove(self, event):
        self.disp.config(text='x: %d, y: %d' % (event.x, event.y), font=('Arial', 15))
        self.STATE['x'], self.STATE['y'] = event.x, event.y
        # self.mainPanel.create_oval(event.x - 5, event.y - 5, event.x + 5, event.y + 5, width=4, outline="lightGreen")

    def mouseClick(self, event):
        if self.STATE['click'] == 1:
            x, y = event.x, event.y
            self.points.append((x, y))
            self.text.insert('insert','(' + str(x) + ', ' + str(y) + ')   ')
            if len(self.points) % 4 == 0:
                self.text.insert('insert', '\n')

    def stepBack(self):
        self.points.pop()

    def clearPoints(self):
        self.points = []
        self.text.delete('1.0', 'end')

    def saveFile(self):
        with open('label.txt', 'a') as f:
            f.write(self.text.get(1.0, END))
        self.text.delete('1.0', 'end')
        self.points = []

    def warn(self):
        messagebox.showwarning(title='无效视频路劲', message='请输入正确的视频路径')

if __name__ == '__main__':
    root = Tk()
    tool = LabelTool(root)
    root.mainloop()