Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import tkinter as tk
- from tkinter import messagebox, filedialog
- import tkinter.ttk as ttk
- from PIL import Image, ImageTk
- import json
- from shapely.geometry import Polygon
- from screeninfo import get_monitors
- import cv2
- class MarkupApp(tk.Tk):
- def __init__(self):
- tk.Tk.__init__(self)
- self.canvas = None
- self.done_btn = None
- self.zone_label = None
- self.clear_btn = None
- self.load_btn = None
- self.cnt_zones = None
- self.tk_img = None
- # self.img_name = "preview_img.png"
- self.img_name = None
- self.preview_name = "preview.png"
- self.colors = ["black", "red", "green", "blue", "cyan", "yellow", "magenta"]
- self.curr_colors = self.colors.copy()
- self.height = None
- self.width = None
- self.click_num = 0
- self.x1 = self.y1 = None
- # self.start_x = self.start_y = None
- self.zones = []
- self.curr_zone = []
- # для удаления
- self.obj_id = []
- self.the_last_clr = None
- # для масштабирования
- self.ratio = 1
- self.init_gui()
- def init_gui(self):
- self.create_canvas()
- self.iconbitmap('icon.ico')
- self.title('Разметка световых зон')
- self.resizable(width=False, height=False)
- self.load_btn = ttk.Button(self, text="Загрузить", command=self.load_image)
- self.load_btn.grid(row=2, column=0, pady=10, ipady=3, ipadx=3)
- self.zone_label = ttk.Label(self, text="Текущее количество зон:")
- self.zone_label.grid(row=0, column=3, sticky='s')
- self.cnt_zones = ttk.Label(self, text='0')
- self.cnt_zones.grid(row=1, column=3, sticky='n')
- self.done_btn = ttk.Button(self, text="Завершить", command=self.save_zone)
- self.done_btn.grid(row=2, column=1, pady=10, ipady=3, ipadx=3)
- # для удаления
- self.done_btn = ttk.Button(self, text="Назад", command=self.back)
- self.done_btn.grid(row=2, column=2, pady=10, ipady=3, ipadx=3)
- def load_image(self):
- filetypes = (
- ('Изображения', '*.jpg; *.jpeg; *.png'),
- ('Все файлы', '*.*')
- )
- self.img_name = filedialog.askopenfilename(title="Выбор изображения для разметки", filetypes=filetypes)
- if self.img_name:
- img = Image.open(self.img_name)
- self.canvas.delete("all")
- self.zones = []
- self.curr_zone = []
- self.x1 = None
- self.y1 = None
- self.click_num = 0
- self.tk_img = ImageTk.PhotoImage(img)
- self.height = self.tk_img.height()
- self.width = self.tk_img.width()
- w_mon = get_monitors()[0].width
- h_mon = get_monitors()[0].height
- if w_mon * 4 / 5.0 < self.width or h_mon * 4 / 5.0 < self.height:
- if self.width > self.height:
- self.ratio = (w_mon * 4 / 5.0) / self.width
- else:
- self.ratio = (h_mon * 4 / 5.0) / self.height
- self.tk_img = ImageTk.PhotoImage((img.resize((int(self.ratio * self.width),
- int(self.ratio * self.height)))))
- self.canvas.config(width=int(self.ratio * self.width), height=int(self.ratio * self.height))
- self.canvas.create_image(0, 0, anchor="nw", image=self.tk_img)
- else:
- self.canvas.config(width=self.width, height=self.height)
- self.canvas.create_image(0, 0, anchor="nw", image=self.tk_img)
- def save_zone(self):
- if len(self.zones) != 0:
- img = cv2.imread(self.img_name)
- depth_net = cv2.dnn.readNet("model-f6b98070.onnx")
- blob = cv2.dnn.blobFromImage(img, 1 / 255., (384, 384), (123.675, 116.28, 103.53), True, False)
- depth_net.setInput(blob)
- depth_map = depth_net.forward()
- depth_map = depth_map[0, :, :]
- depth_map = cv2.resize(depth_map, (self.width, self.height))
- depth_map = cv2.normalize(depth_map, None, 0, 1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
- res_zones = {f"zone_{i}": coords for i, coords in enumerate(self.zones)}
- res_zones["depth_map"] = depth_map.tolist()
- with open('markup_data.json', 'w') as f:
- json.dump(res_zones, f)
- messagebox.showinfo("Завершение!", "Разметка успешно завершена. Вы можете закрыть приложение.")
- else:
- messagebox.showwarning("Предупреждение!", "Для начала выделите хотя бы одну зону.")
- # messagebox.showinfo("Завершение!", "Разметка успешно завершена. Вы можете закрыть приложение.")
- # res_zones = {f"zone_{i}": coords for i, coords in enumerate(self.zones)}
- # with open('markup_data.json', 'w') as f:
- # json.dump(res_zones, f)
- # else:
- # messagebox.showwarning("Предупреждение!", "Для начала выделите хотя бы одну зону.")
- def count_zones(self):
- self.cnt_zones.config(text='{}'.format(len(self.zones)))
- def create_canvas(self):
- if self.img_name:
- img = Image.open(self.img_name)
- else:
- img = Image.open(self.preview_name)
- self.tk_img = ImageTk.PhotoImage(img)
- self.height = self.tk_img.height()
- self.width = self.tk_img.width()
- self.canvas = tk.Canvas(self, width=self.width, height=self.height, background="white")
- self.canvas.grid(row=0, columnspan=3, rowspan=2, padx=10)
- self.canvas.bind('<Button-1>', self.draw_line)
- self.canvas.create_image(0, 0, anchor="nw", image=self.tk_img)
- def draw_line(self, event):
- if not self.img_name:
- return
- if self.click_num == 0:
- self.curr_zone = []
- self.x1 = event.x
- self.y1 = event.y
- if len(self.curr_colors) == 0:
- self.curr_colors = self.colors.copy()
- self.obj_id.append(self.canvas.create_oval(self.x1 - 2, self.y1 - 2, self.x1 + 2, self.y1 + 2, width=0,
- fill=self.curr_colors[-1]))
- self.click_num = 1
- self.curr_zone.append({"x": self.x1 / self.ratio, "y": self.y1 / self.ratio})
- else:
- x2 = event.x
- y2 = event.y
- if self.click_num > 2 and \
- (abs(self.curr_zone[0]['x'] * self.ratio - x2) < 10 and
- abs(self.curr_zone[0]['y'] * self.ratio - y2) < 10):
- self.obj_id.append(self.canvas.create_line(self.x1,
- self.y1,
- self.curr_zone[0]['x'] * self.ratio,
- self.curr_zone[0]['y'] * self.ratio,
- fill=self.curr_colors[-1],
- width=2))
- coord = self.make_coord()
- self.obj_id.append(self.canvas.create_text(coord.x,
- coord.y,
- text=str(len(self.zones) + 1),
- fill=self.curr_colors[-1],
- font=('Helvetica 18 bold')))
- self.click_num = 0
- self.the_last_clr = self.curr_colors[-1]
- self.curr_colors.pop()
- self.zones.append(self.curr_zone)
- self.count_zones()
- else:
- self.curr_zone.append({"x": x2 / self.ratio, "y": y2 / self.ratio})
- self.obj_id.append(self.canvas.create_oval(x2 - 2, y2 - 2, x2 + 2, y2 + 2, width=0,
- fill=self.curr_colors[-1]))
- self.obj_id.append(self.canvas.create_line(self.x1, self.y1, x2, y2,
- fill=self.curr_colors[-1], width=2))
- self.x1 = x2
- self.y1 = y2
- self.click_num += 1
- def make_coord(self):
- ar = []
- for i in self.curr_zone:
- ar.append([i['x'] * self.ratio, i['y'] * self.ratio])
- return Polygon(ar).centroid
- def back(self):
- if self.click_num == 0 and len(self.curr_zone) != 0:
- for i in range(2):
- self.canvas.delete(self.obj_id[-1])
- self.obj_id.pop()
- self.zones.pop()
- self.click_num = len(self.curr_zone)
- self.curr_colors.append(self.the_last_clr)
- self.count_zones()
- elif len(self.curr_zone) > 1:
- for i in range(2):
- self.canvas.delete(self.obj_id[-1])
- self.obj_id.pop()
- self.click_num -= 1
- self.curr_zone.pop()
- self.x1 = self.curr_zone[-1]['x'] * self.ratio
- self.y1 = self.curr_zone[-1]['y'] * self.ratio
- elif len(self.curr_zone) == 1:
- self.canvas.delete(self.obj_id[-1])
- self.obj_id.pop()
- self.click_num -= 1
- self.curr_zone.pop()
- elif len(self.zones) and len(self.curr_zone) == 0:
- self.curr_zone = self.zones[-1]
- for i in range(len(self.colors)):
- if self.the_last_clr == self.colors[i]:
- if i == len(self.colors) - 1:
- self.the_last_clr = self.colors[0]
- else:
- self.the_last_clr = self.colors[i + 1]
- break
- self.x1 = self.curr_zone[-1]['x'] * self.ratio
- self.y1 = self.curr_zone[-1]['y'] * self.ratio
- # self.start_x = self.curr_zone[0]['x']
- # self.start_y = self.curr_zone[0]['y']
- self.back()
- if __name__ == '__main__':
- app = MarkupApp()
- app.mainloop()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement