Advertisement
Prin_Di

new_gui

May 13th, 2022
517
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import tkinter as tk
  2. from tkinter import messagebox, filedialog
  3. import tkinter.ttk as ttk
  4. from PIL import Image, ImageTk
  5. import json
  6. from shapely.geometry import Polygon
  7. from screeninfo import get_monitors
  8. import cv2
  9.  
  10.  
  11. class MarkupApp(tk.Tk):
  12.     def __init__(self):
  13.         tk.Tk.__init__(self)
  14.         self.canvas = None
  15.         self.done_btn = None
  16.         self.zone_label = None
  17.         self.clear_btn = None
  18.         self.load_btn = None
  19.         self.cnt_zones = None
  20.  
  21.         self.tk_img = None
  22.         # self.img_name = "preview_img.png"
  23.         self.img_name = None
  24.         self.preview_name = "preview.png"
  25.         self.colors = ["black", "red", "green", "blue", "cyan", "yellow", "magenta"]
  26.         self.curr_colors = self.colors.copy()
  27.  
  28.         self.height = None
  29.         self.width = None
  30.  
  31.         self.click_num = 0
  32.         self.x1 = self.y1 = None
  33.         # self.start_x = self.start_y = None
  34.  
  35.         self.zones = []
  36.         self.curr_zone = []
  37.  
  38.         # для удаления
  39.         self.obj_id = []
  40.         self.the_last_clr = None
  41.         # для масштабирования
  42.         self.ratio = 1
  43.  
  44.         self.init_gui()
  45.  
  46.     def init_gui(self):
  47.         self.create_canvas()
  48.         self.iconbitmap('icon.ico')
  49.         self.title('Разметка световых зон')
  50.         self.resizable(width=False, height=False)
  51.         self.load_btn = ttk.Button(self, text="Загрузить", command=self.load_image)
  52.         self.load_btn.grid(row=2, column=0, pady=10, ipady=3, ipadx=3)
  53.  
  54.         self.zone_label = ttk.Label(self, text="Текущее количество зон:")
  55.         self.zone_label.grid(row=0, column=3, sticky='s')
  56.  
  57.         self.cnt_zones = ttk.Label(self, text='0')
  58.         self.cnt_zones.grid(row=1, column=3, sticky='n')
  59.  
  60.         self.done_btn = ttk.Button(self, text="Завершить", command=self.save_zone)
  61.         self.done_btn.grid(row=2, column=1, pady=10, ipady=3, ipadx=3)
  62.  
  63.         # для удаления
  64.         self.done_btn = ttk.Button(self, text="Назад", command=self.back)
  65.         self.done_btn.grid(row=2, column=2, pady=10, ipady=3, ipadx=3)
  66.  
  67.     def load_image(self):
  68.         filetypes = (
  69.             ('Изображения', '*.jpg; *.jpeg; *.png'),
  70.             ('Все файлы', '*.*')
  71.         )
  72.         self.img_name = filedialog.askopenfilename(title="Выбор изображения для разметки", filetypes=filetypes)
  73.         if self.img_name:
  74.             img = Image.open(self.img_name)
  75.             self.canvas.delete("all")
  76.             self.zones = []
  77.             self.curr_zone = []
  78.             self.x1 = None
  79.             self.y1 = None
  80.             self.click_num = 0
  81.             self.tk_img = ImageTk.PhotoImage(img)
  82.             self.height = self.tk_img.height()
  83.             self.width = self.tk_img.width()
  84.             w_mon = get_monitors()[0].width
  85.             h_mon = get_monitors()[0].height
  86.             if w_mon * 4 / 5.0 < self.width or h_mon * 4 / 5.0 < self.height:
  87.                 if self.width > self.height:
  88.                     self.ratio = (w_mon * 4 / 5.0) / self.width
  89.                 else:
  90.                     self.ratio = (h_mon * 4 / 5.0) / self.height
  91.                 self.tk_img = ImageTk.PhotoImage((img.resize((int(self.ratio * self.width),
  92.                                                               int(self.ratio * self.height)))))
  93.                 self.canvas.config(width=int(self.ratio * self.width), height=int(self.ratio * self.height))
  94.                 self.canvas.create_image(0, 0, anchor="nw", image=self.tk_img)
  95.             else:
  96.                 self.canvas.config(width=self.width, height=self.height)
  97.                 self.canvas.create_image(0, 0, anchor="nw", image=self.tk_img)
  98.  
  99.     def save_zone(self):
  100.         if len(self.zones) != 0:
  101.             img = cv2.imread(self.img_name)
  102.             depth_net = cv2.dnn.readNet("model-f6b98070.onnx")
  103.             blob = cv2.dnn.blobFromImage(img, 1 / 255., (384, 384), (123.675, 116.28, 103.53), True, False)
  104.             depth_net.setInput(blob)
  105.             depth_map = depth_net.forward()
  106.             depth_map = depth_map[0, :, :]
  107.             depth_map = cv2.resize(depth_map, (self.width, self.height))
  108.             depth_map = cv2.normalize(depth_map, None, 0, 1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
  109.             res_zones = {f"zone_{i}": coords for i, coords in enumerate(self.zones)}
  110.             res_zones["depth_map"] = depth_map.tolist()
  111.             with open('markup_data.json', 'w') as f:
  112.                 json.dump(res_zones, f)
  113.             messagebox.showinfo("Завершение!", "Разметка успешно завершена. Вы можете закрыть приложение.")
  114.         else:
  115.             messagebox.showwarning("Предупреждение!", "Для начала выделите хотя бы одну зону.")
  116.         #     messagebox.showinfo("Завершение!", "Разметка успешно завершена. Вы можете закрыть приложение.")
  117.         #     res_zones = {f"zone_{i}": coords for i, coords in enumerate(self.zones)}
  118.         #     with open('markup_data.json', 'w') as f:
  119.         #         json.dump(res_zones, f)
  120.         # else:
  121.         #     messagebox.showwarning("Предупреждение!", "Для начала выделите хотя бы одну зону.")
  122.  
  123.     def count_zones(self):
  124.         self.cnt_zones.config(text='{}'.format(len(self.zones)))
  125.  
  126.     def create_canvas(self):
  127.         if self.img_name:
  128.             img = Image.open(self.img_name)
  129.         else:
  130.             img = Image.open(self.preview_name)
  131.         self.tk_img = ImageTk.PhotoImage(img)
  132.         self.height = self.tk_img.height()
  133.         self.width = self.tk_img.width()
  134.         self.canvas = tk.Canvas(self, width=self.width, height=self.height, background="white")
  135.         self.canvas.grid(row=0, columnspan=3, rowspan=2, padx=10)
  136.         self.canvas.bind('<Button-1>', self.draw_line)
  137.         self.canvas.create_image(0, 0, anchor="nw", image=self.tk_img)
  138.  
  139.     def draw_line(self, event):
  140.         if not self.img_name:
  141.             return
  142.         if self.click_num == 0:
  143.             self.curr_zone = []
  144.             self.x1 = event.x
  145.             self.y1 = event.y
  146.             if len(self.curr_colors) == 0:
  147.                 self.curr_colors = self.colors.copy()
  148.             self.obj_id.append(self.canvas.create_oval(self.x1 - 2, self.y1 - 2, self.x1 + 2, self.y1 + 2, width=0,
  149.                                                        fill=self.curr_colors[-1]))
  150.             self.click_num = 1
  151.             self.curr_zone.append({"x": self.x1 / self.ratio, "y": self.y1 / self.ratio})
  152.         else:
  153.             x2 = event.x
  154.             y2 = event.y
  155.             if self.click_num > 2 and \
  156.                     (abs(self.curr_zone[0]['x'] * self.ratio - x2) < 10 and
  157.                      abs(self.curr_zone[0]['y'] * self.ratio - y2) < 10):
  158.                 self.obj_id.append(self.canvas.create_line(self.x1,
  159.                                                            self.y1,
  160.                                                            self.curr_zone[0]['x'] * self.ratio,
  161.                                                            self.curr_zone[0]['y'] * self.ratio,
  162.                                                            fill=self.curr_colors[-1],
  163.                                                            width=2))
  164.  
  165.                 coord = self.make_coord()
  166.                 self.obj_id.append(self.canvas.create_text(coord.x,
  167.                                                            coord.y,
  168.                                                            text=str(len(self.zones) + 1),
  169.                                                            fill=self.curr_colors[-1],
  170.                                                            font=('Helvetica 18 bold')))
  171.                 self.click_num = 0
  172.                 self.the_last_clr = self.curr_colors[-1]
  173.                 self.curr_colors.pop()
  174.                 self.zones.append(self.curr_zone)
  175.                 self.count_zones()
  176.             else:
  177.                 self.curr_zone.append({"x": x2 / self.ratio, "y": y2 / self.ratio})
  178.                 self.obj_id.append(self.canvas.create_oval(x2 - 2, y2 - 2, x2 + 2, y2 + 2, width=0,
  179.                                                            fill=self.curr_colors[-1]))
  180.                 self.obj_id.append(self.canvas.create_line(self.x1, self.y1, x2, y2,
  181.                                                            fill=self.curr_colors[-1], width=2))
  182.                 self.x1 = x2
  183.                 self.y1 = y2
  184.                 self.click_num += 1
  185.  
  186.     def make_coord(self):
  187.         ar = []
  188.         for i in self.curr_zone:
  189.             ar.append([i['x'] * self.ratio, i['y'] * self.ratio])
  190.         return Polygon(ar).centroid
  191.  
  192.     def back(self):
  193.         if self.click_num == 0 and len(self.curr_zone) != 0:
  194.             for i in range(2):
  195.                 self.canvas.delete(self.obj_id[-1])
  196.                 self.obj_id.pop()
  197.             self.zones.pop()
  198.             self.click_num = len(self.curr_zone)
  199.             self.curr_colors.append(self.the_last_clr)
  200.             self.count_zones()
  201.         elif len(self.curr_zone) > 1:
  202.             for i in range(2):
  203.                 self.canvas.delete(self.obj_id[-1])
  204.                 self.obj_id.pop()
  205.             self.click_num -= 1
  206.             self.curr_zone.pop()
  207.             self.x1 = self.curr_zone[-1]['x'] * self.ratio
  208.             self.y1 = self.curr_zone[-1]['y'] * self.ratio
  209.         elif len(self.curr_zone) == 1:
  210.             self.canvas.delete(self.obj_id[-1])
  211.             self.obj_id.pop()
  212.             self.click_num -= 1
  213.             self.curr_zone.pop()
  214.  
  215.         elif len(self.zones) and len(self.curr_zone) == 0:
  216.             self.curr_zone = self.zones[-1]
  217.             for i in range(len(self.colors)):
  218.                 if self.the_last_clr == self.colors[i]:
  219.                     if i == len(self.colors) - 1:
  220.                         self.the_last_clr = self.colors[0]
  221.                     else:
  222.                         self.the_last_clr = self.colors[i + 1]
  223.                     break
  224.             self.x1 = self.curr_zone[-1]['x'] * self.ratio
  225.             self.y1 = self.curr_zone[-1]['y'] * self.ratio
  226.             # self.start_x = self.curr_zone[0]['x']
  227.             # self.start_y = self.curr_zone[0]['y']
  228.             self.back()
  229.  
  230.  
  231. if __name__ == '__main__':
  232.     app = MarkupApp()
  233.     app.mainloop()
  234.  
Advertisement
Advertisement
Advertisement
RAW Paste Data Copied
Advertisement