forked from ultralytics/yolov5
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathlabel_tool.py
143 lines (128 loc) · 5.2 KB
/
label_tool.py
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
137
138
139
140
141
142
143
from pathlib import Path
import argparse
import numpy as np
import cv2 as cv
from utils.general import increment_path
from myutils.general import loadcsv
IMG_TYPES = [".jpg", ".png"]
save_dir = Path(increment_path("runs/label/exp", mkdir=True))
cv.namedWindow('frame', cv.WINDOW_NORMAL)
cv.resizeWindow('frame', (720, 1080))
ori_img = None
painted = None
labels = None
labeling = True
pt_front, pt_back = None, None
pressed = None
show_label = True
def repaint(img, labels):
global show_label
if show_label:
for x, y in labels:
img = cv.circle(img, (x, y), 4, (255, 0, 0), -1)
return img
def mouse_clicked_callback(event, x, y, flags, params):
def inside_rect(x, y, left_up, right_bot):
return right_bot[0] > x > left_up[0] and right_bot[1] > y > left_up[1]
global ori_img, painted, labels, pt_front, pt_back, labeling
if event == cv.EVENT_LBUTTONDBLCLK:
if labeling:
labels = np.vstack((labels, (x, y)))
print(x, y, "added to labels")
else:
if pt_front is None:
pt_front = np.array((x, y), dtype=int)
else:
pt_back = np.array((x, y), dtype=int)
pt1 = np.array((min(pt_front[0], pt_back[0]), min(
pt_front[1], pt_back[1])), dtype=int)
pt2 = np.array((max(pt_front[0], pt_back[0]), max(
pt_front[1], pt_back[1])), dtype=int)
n_old = labels.shape[0]
xy_pts = labels[:, [0, 1]]
labels = np.delete(labels, np.all(
((xy_pts > pt1) & (xy_pts < pt2)), axis=1), axis=0)
n_new = labels.shape[0]
print(f"{n_old - n_new} labels from {pt1}, {pt2} removed. ")
pt_front = None
painted = repaint(ori_img.copy(), labels)
if event == cv.EVENT_MOUSEMOVE and not pt_front is None and not labeling:
pt_back = np.array((x, y), dtype=int)
pt1 = (min(pt_front[0], pt_back[0]), min(pt_front[1], pt_back[1]))
pt2 = (max(pt_front[0], pt_back[0]), max(pt_front[1], pt_back[1]))
painted = cv.rectangle(ori_img.copy(), pt1, pt2, (0, 255, 0), 3)
painted = repaint(painted, labels)
cv.setMouseCallback('frame', mouse_clicked_callback)
def main(opt):
global ori_img, painted, labels, labeling, show_label, pt_front
img_list = [p for p in opt.dir.iterdir() if p.suffix.lower() in IMG_TYPES]
n_imgs = len(img_list)
idx = 0
ori_img = cv.imread(str(img_list[idx]))
assert 'labels' in opt, 'Must be labeled before'
label_p = opt.labels / f"{img_list[idx].stem}.csv"
labels = loadcsv(label_p)
painted = repaint(ori_img.copy(), labels)
print("Controls:")
print(" a: previous")
print(" d: next")
print(" c: with/without labels")
print(" e: switch mode between labeling & erasing")
print(" q: quit")
while True:
cv.imshow("frame", painted)
key = cv.waitKey(10)
if (key == ord("a") and idx > 0) or (key == ord("d") and idx < n_imgs - 1):
np.savetxt(str(save_dir / label_p.name),
labels, fmt='%d', delimiter=',')
if key == ord('a') and idx > 0:
idx = idx - 1
elif key == ord("d") and idx < n_imgs - 1:
idx = idx + 1
print(str(img_list[idx]))
ori_img = cv.imread(str(img_list[idx]))
label_p = opt.labels / f"{img_list[idx].stem}.csv"
if (save_dir / label_p.name).is_file():
label_p = save_dir / f"{img_list[idx].stem}.csv"
labels = loadcsv(label_p)
painted = repaint(ori_img.copy(), labels)
elif key == ord("e"):
labeling = not labeling
s = "label mode" if labeling else "erase mode"
pt_front = None
print("Switching to", s)
elif key == ord('c'):
show_label = not show_label
painted = repaint(ori_img.copy(), labels)
elif key == ord("q"):
np.savetxt(str(save_dir / label_p.name),
labels, fmt='%d', delimiter=',')
break
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-d", "--dir",
help="Directory of images")
parser.add_argument("-l", "--labels",
help="Directory of all labels")
opt = parser.parse_args()
# opt.dir = Path.cwd() / "dataset/raw/test/test_public/d"
# opt.labels = Path.cwd() / "uploads/0607-2324/upload"
# test_dir = Path.cwd() / "dataset/raw/test"
# if not (Path.cwd() / (opt.dir)).is_dir():
# if opt.dir == "public":
# opt.dir = test_dir / "test_public_all"
# elif opt.dir == "i":
# opt.dir = test_dir / 'test_public/i'
# elif opt.dir == "d":
# opt.dir = test_dir / 'test_public/d'
# elif opt.dir == "private":
# opt.dir = test_dir / "test_private"
# elif opt.dir == "all":
# opt.dir = test_dir / "test_all"
# else:
# print("Dir not found!")
# exit(-1)
# print("Redirect to", opt.dir)
opt.dir = Path(opt.dir)
opt.labels = Path(opt.labels)
main(opt)