cv-e2e-playground/annotate.py at main · macrobomaster/cv-e2e-playground · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
import os

import cv2
import glob
from pathlib import Path

from main import BASE_PATH

if (stage := int(os.environ.get("STAGE", 0))) == 0:
  frame_files = glob.glob(str(BASE_PATH / "data/*.png"))
  print(f"there are {len(frame_files)} frames")

  oframe, frame = None, None
  click_pos = None

  def click_handler(event, x, y, flags, param):
    global oframe, frame, click_pos
    if event == cv2.EVENT_LBUTTONDOWN:
      frame = oframe.copy()  # type: ignore
      cv2.circle(frame, (x, y), 3, (100, 0, 255), -1)
      click_pos = (x, y)

  cv2.namedWindow("preview")
  cv2.setMouseCallback("preview", click_handler)

  i = 0
  flag = False
  while not flag and i < len(frame_files):
    print(f"annotating frame {i} of {len(frame_files)}")
    frame_file = frame_files[i]
    oframe = frame = cv2.imread(frame_file)

    while True:
      cv2.imshow("preview", frame)

      key = cv2.waitKey(1)
      if key == ord("q"):
        flag = True
        break
      elif key == ord("s"):
        with open(Path(frame_file).with_suffix(".txt"), "w") as f:
          f.write(f"0 0 0")
        break
      elif key == ord("a") and click_pos is not None:
        with open(Path(frame_file).with_suffix(".txt"), "w") as f:
          f.write(f"1 {click_pos[0]} {click_pos[1]}")  # type: ignore
        click_pos = None
        break
      elif key == ord("d"):
        i -= 2
        break
    i += 1
elif stage == 1:
  frame_files = glob.glob(str(BASE_PATH / "data/*.png"))
  print(f"there are {len(frame_files)} frames")

  i = 0
  flag = False
  while not flag and i < len(frame_files):
    print(f"annotating color of frame {i} of {len(frame_files)}")
    frame_file = frame_files[i]
    frame = cv2.imread(frame_file)
    cv2.imshow("preview", frame)

    while True:
      key = cv2.waitKey(0)
      if key == ord("q"):
        flag = True
        break
      elif key == ord("s"):
        with open(Path(frame_file).with_suffix(".txt"), "r") as f:
          line = f.readline().split(" ")
          detected, x, y = int(line[0]), int(line[1]), int(line[2])
        with open(Path(frame_file).with_suffix(".txt"), "w") as f:
          f.write(f"{detected} {x} {y} -1")
        break
      elif key == ord("r"):
        with open(Path(frame_file).with_suffix(".txt"), "r") as f:
          line = f.readline().split(" ")
          detected, x, y = int(line[0]), int(line[1]), int(line[2])
        with open(Path(frame_file).with_suffix(".txt"), "w") as f:
          f.write(f"{detected} {x} {y} 0")
        break
      elif key == ord("b"):
        with open(Path(frame_file).with_suffix(".txt"), "r") as f:
          line = f.readline().split(" ")
          detected, x, y = int(line[0]), int(line[1]), int(line[2])
        with open(Path(frame_file).with_suffix(".txt"), "w") as f:
          f.write(f"{detected} {x} {y} 1")
        break
      elif key == ord("d"):
        i -= 2
        break
    i += 1
elif stage == 2:
  frame_files = glob.glob(str(BASE_PATH / "data/*.png"))
  print(f"there are {len(frame_files)} frames")

  i = 0
  flag = False
  while not flag and i < len(frame_files):
    print(f"annotating number of frame {i} of {len(frame_files)}")
    frame_file = frame_files[i]
    frame = cv2.imread(frame_file)
    cv2.imshow("preview", frame)

    while True:
      key = cv2.waitKey(0)
      if key == ord("q"):
        flag = True
        break
      elif key == ord("s"):
        with open(Path(frame_file).with_suffix(".txt"), "r") as f:
          line = f.readline().split(" ")
          detected, x, y, color = int(line[0]), int(line[1]), int(line[2]), int(line[3])
        with open(Path(frame_file).with_suffix(".txt"), "w") as f:
          f.write(f"{detected} {x} {y} {color} -1")
        break
      elif key == ord("1"):
        with open(Path(frame_file).with_suffix(".txt"), "r") as f:
          line = f.readline().split(" ")
          detected, x, y, color = int(line[0]), int(line[1]), int(line[2]), int(line[3])
        with open(Path(frame_file).with_suffix(".txt"), "w") as f:
          f.write(f"{detected} {x} {y} {color} 1")
        break
      elif key == ord("3"):
        with open(Path(frame_file).with_suffix(".txt"), "r") as f:
          line = f.readline().split(" ")
          detected, x, y, color = int(line[0]), int(line[1]), int(line[2]), int(line[3])
        with open(Path(frame_file).with_suffix(".txt"), "w") as f:
          f.write(f"{detected} {x} {y} {color} 3")
        break
      elif key == ord("d"):
        i -= 2
        break
    i += 1