def __init__(self, size, margin, weights, device, select_largest, images):
print("Loading network.....")
current = time.time()
self.device = device
self.mtcnn = MTCNN(image_size=size,
margin=margin,
keep_all=False,
device=device,
select_largest=True)
self.model = InceptionResnetV1(pretrained='vggface2', device=device)
state_dict = torch.load(weights, map_location="cpu")
self.model.load_state_dict(state_dict)
self.images = images
elapse = time.time() - current
print("Network successfully loaded, process time: %.2fs" % elapse)
# eval mode
self.model.eval()
def extract_faces(videopath):
mtcnn = MTCNN(device='cuda', image_size=299, margin=0.3 / 1.3 * 299).eval()
reader = cv2.VideoCapture(videopath)
frame_list = []
frame_num = 0
while reader.isOpened():
success, frame = reader.read()
if not success:
break
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame_list.append(image)
frame_num += 1
reader.release()
# print(len(frame_list[:10]))
faces = []
with torch.no_grad():
face = mtcnn(frame_list[:2])
faces.extend(face)
faces = [f for f in faces if f is not None]
return faces[0]
def __init__(self):
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.cnn = MTCNN(device=self.device, keep_all=True)
self.cnn0 = MTCNN(device=self.device)
self.rnn1 = InceptionResnetV1(device=self.device, pretrained='vggface2').eval()
def __init__(self, master, opt):
self.parent = master
self.parent.title("Semi Automatic Image Annotation Tool")
self.frame = Frame(self.parent)
self.frame.pack(fill=BOTH, expand=1)
self.parent.resizable(width=False, height=False)
# setup yolov5 model
self.opt = opt
opt.augment = True
opt.update = False
opt.agnostic_nms = False
# initialize
set_logging()
device = select_device(opt.device)
self.device = device
# load yolo5 object detection model
model = attempt_load(opt.weights,
map_location=device) # load FP32 model
self.img_size = check_img_size(opt.img_size,
s=model.stride.max()) # check img_size
half = device.type != 'cpu' # half precision only supported on CUDA
self.half = half
if half:
model.half() # to FP16
self.object_model = model
# get object detection names
self.names = ['person']
# load face detect model
self.face_model = MTCNN(keep_all=True,
margin=opt.face_margin,
device=device)
self.face_landmarks = opt.face_landmarks
# Initialize class variables
self.img = None
self.tkimg = None
self.imageDir = ''
self.imageDirPathBuffer = ''
self.imageList = []
self.imageTotal = 0
self.imageCur = 0
self.cur = 0
self.bboxIdList = []
self.bboxList = []
self.bboxPointList = []
self.o1 = None
self.o2 = None
self.o3 = None
self.o4 = None
self.bboxId = None
self.currLabel = None
self.editbboxId = None
self.currBboxColor = None
self.zoomImgId = None
self.zoomImg = None
self.zoomImgCrop = None
self.tkZoomImg = None
self.hl = None
self.vl = None
self.editPointId = None
self.filename = None
self.filenameBuffer = None
self.objectLabelList = []
self.EDIT = False
# record image names for saving previous results
self.annoList = {}
# initialize mouse state
self.STATE = {'x': 0, 'y': 0}
self.STATE_COCO = {'click': 0}
# initialize annotation file
self.anno_filename = 'annotations.csv'
self.annotation_file = open('annotations/' + self.anno_filename, 'w+')
self.annotation_file.write("")
self.annotation_file.close()
# ------------------ GUI ---------------------
# Control Panel
self.ctrlPanel = Frame(self.frame)
self.ctrlPanel.grid(row=0, column=0, sticky=W + N)
self.openBtn = Button(self.ctrlPanel,
text='Open',
command=self.open_image)
self.openBtn.pack(fill=X, side=TOP)
self.openDirBtn = Button(self.ctrlPanel,
text='Open Dir',
command=self.open_image_dir)
self.openDirBtn.pack(fill=X, side=TOP)
self.nextBtn = Button(self.ctrlPanel,
text='Next -->',
command=self.open_next)
self.nextBtn.pack(fill=X, side=TOP)
self.previousBtn = Button(self.ctrlPanel,
text='<-- Previous',
command=self.open_previous)
self.previousBtn.pack(fill=X, side=TOP)
self.saveBtn = Button(self.ctrlPanel, text='Save', command=self.save)
self.saveBtn.pack(fill=X, side=TOP)
self.semiAutoBtn = Button(self.ctrlPanel,
text="Show Suggestions",
command=self.automate)
self.semiAutoBtn.pack(fill=X, side=TOP)
self.disp = Label(self.ctrlPanel, text='Coordinates:')
self.disp.pack(fill=X, side=TOP)
self.mb = Menubutton(self.ctrlPanel,
text="COCO Classes for Suggestions",
relief=RAISED)
self.mb.pack(fill=X, side=TOP)
self.mb.menu = Menu(self.mb, tearoff=0)
self.mb["menu"] = self.mb.menu
self.addCocoBtn = Button(self.ctrlPanel,
text="+",
command=self.add_labels_coco)
self.addCocoBtn.pack(fill=X, side=TOP)
self.zoomPanelLabel = Label(self.ctrlPanel,
text="Precision View Panel")
self.zoomPanelLabel.pack(fill=X, side=TOP)
self.zoomcanvas = Canvas(self.ctrlPanel, width=150, height=150)
self.zoomcanvas.pack(fill=X, side=TOP, anchor='center')
# Image Editing Region
self.canvas = Canvas(self.frame,
width=self.img_size,
height=self.img_size)
self.canvas.grid(row=0, column=1, sticky=W + N)
self.canvas.bind("<Button-1>", self.mouse_click)
self.canvas.bind("<Motion>", self.mouse_move, "+")
self.canvas.bind("<B1-Motion>", self.mouse_drag)
self.canvas.bind("<ButtonRelease-1>", self.mouse_release)
self.parent.bind("<Key-Left>", self.open_previous)
self.parent.bind("<Key-Right>", self.open_next)
self.parent.bind("Escape", self.cancel_bbox)
# Labels and Bounding Box Lists Panel
self.listPanel = Frame(self.frame)
self.listPanel.grid(row=0, column=2, sticky=W + N)
self.listBoxNameLabel = Label(self.listPanel,
text="List of Objects").pack(fill=X,
side=TOP)
self.objectListBox = Listbox(self.listPanel, width=40)
self.objectListBox.pack(fill=X, side=TOP)
self.delObjectBtn = Button(self.listPanel,
text="Delete",
command=self.del_bbox)
self.delObjectBtn.pack(fill=X, side=TOP)
self.clearAllBtn = Button(self.listPanel,
text="Clear All",
command=self.clear_bbox)
self.clearAllBtn.pack(fill=X, side=TOP)
self.classesNameLabel = Label(self.listPanel,
text="Classes").pack(fill=X, side=TOP)
self.textBox = Entry(self.listPanel, text="Enter label")
self.textBox.pack(fill=X, side=TOP)
self.addLabelBtn = Button(self.listPanel,
text="+",
command=self.add_label).pack(fill=X,
side=TOP)
self.delLabelBtn = Button(self.listPanel,
text="-",
command=self.del_label).pack(fill=X,
side=TOP)
self.labelListBox = Listbox(self.listPanel)
self.labelListBox.pack(fill=X, side=TOP)
for name in self.names + ['face']:
self.labelListBox.insert(END, str(name))
self.cocoLabels = config.labels_to_names.values()
self.cocoIntVars = []
for idxcoco, label_coco in enumerate(self.cocoLabels):
self.cocoIntVars.append(IntVar())
self.mb.menu.add_checkbutton(label=label_coco,
variable=self.cocoIntVars[idxcoco])
# print(self.cocoIntVars)
# STATUS BAR
self.statusBar = Frame(self.frame, width=500)
self.statusBar.grid(row=1, column=1, sticky=W + N)
self.processingLabel = Label(self.statusBar,
text=" ")
self.processingLabel.pack(side="left", fill=X)
self.imageIdxLabel = Label(self.statusBar,
text=" ")
self.imageIdxLabel.pack(side="right", fill=X)
class MainGUI:
def __init__(self, master, opt):
self.parent = master
self.parent.title("Semi Automatic Image Annotation Tool")
self.frame = Frame(self.parent)
self.frame.pack(fill=BOTH, expand=1)
self.parent.resizable(width=False, height=False)
# setup yolov5 model
self.opt = opt
opt.augment = True
opt.update = False
opt.agnostic_nms = False
# initialize
set_logging()
device = select_device(opt.device)
self.device = device
# load yolo5 object detection model
model = attempt_load(opt.weights,
map_location=device) # load FP32 model
self.img_size = check_img_size(opt.img_size,
s=model.stride.max()) # check img_size
half = device.type != 'cpu' # half precision only supported on CUDA
self.half = half
if half:
model.half() # to FP16
self.object_model = model
# get object detection names
self.names = ['person']
# load face detect model
self.face_model = MTCNN(keep_all=True,
margin=opt.face_margin,
device=device)
self.face_landmarks = opt.face_landmarks
# Initialize class variables
self.img = None
self.tkimg = None
self.imageDir = ''
self.imageDirPathBuffer = ''
self.imageList = []
self.imageTotal = 0
self.imageCur = 0
self.cur = 0
self.bboxIdList = []
self.bboxList = []
self.bboxPointList = []
self.o1 = None
self.o2 = None
self.o3 = None
self.o4 = None
self.bboxId = None
self.currLabel = None
self.editbboxId = None
self.currBboxColor = None
self.zoomImgId = None
self.zoomImg = None
self.zoomImgCrop = None
self.tkZoomImg = None
self.hl = None
self.vl = None
self.editPointId = None
self.filename = None
self.filenameBuffer = None
self.objectLabelList = []
self.EDIT = False
# record image names for saving previous results
self.annoList = {}
# initialize mouse state
self.STATE = {'x': 0, 'y': 0}
self.STATE_COCO = {'click': 0}
# initialize annotation file
self.anno_filename = 'annotations.csv'
self.annotation_file = open('annotations/' + self.anno_filename, 'w+')
self.annotation_file.write("")
self.annotation_file.close()
# ------------------ GUI ---------------------
# Control Panel
self.ctrlPanel = Frame(self.frame)
self.ctrlPanel.grid(row=0, column=0, sticky=W + N)
self.openBtn = Button(self.ctrlPanel,
text='Open',
command=self.open_image)
self.openBtn.pack(fill=X, side=TOP)
self.openDirBtn = Button(self.ctrlPanel,
text='Open Dir',
command=self.open_image_dir)
self.openDirBtn.pack(fill=X, side=TOP)
self.nextBtn = Button(self.ctrlPanel,
text='Next -->',
command=self.open_next)
self.nextBtn.pack(fill=X, side=TOP)
self.previousBtn = Button(self.ctrlPanel,
text='<-- Previous',
command=self.open_previous)
self.previousBtn.pack(fill=X, side=TOP)
self.saveBtn = Button(self.ctrlPanel, text='Save', command=self.save)
self.saveBtn.pack(fill=X, side=TOP)
self.semiAutoBtn = Button(self.ctrlPanel,
text="Show Suggestions",
command=self.automate)
self.semiAutoBtn.pack(fill=X, side=TOP)
self.disp = Label(self.ctrlPanel, text='Coordinates:')
self.disp.pack(fill=X, side=TOP)
self.mb = Menubutton(self.ctrlPanel,
text="COCO Classes for Suggestions",
relief=RAISED)
self.mb.pack(fill=X, side=TOP)
self.mb.menu = Menu(self.mb, tearoff=0)
self.mb["menu"] = self.mb.menu
self.addCocoBtn = Button(self.ctrlPanel,
text="+",
command=self.add_labels_coco)
self.addCocoBtn.pack(fill=X, side=TOP)
self.zoomPanelLabel = Label(self.ctrlPanel,
text="Precision View Panel")
self.zoomPanelLabel.pack(fill=X, side=TOP)
self.zoomcanvas = Canvas(self.ctrlPanel, width=150, height=150)
self.zoomcanvas.pack(fill=X, side=TOP, anchor='center')
# Image Editing Region
self.canvas = Canvas(self.frame,
width=self.img_size,
height=self.img_size)
self.canvas.grid(row=0, column=1, sticky=W + N)
self.canvas.bind("<Button-1>", self.mouse_click)
self.canvas.bind("<Motion>", self.mouse_move, "+")
self.canvas.bind("<B1-Motion>", self.mouse_drag)
self.canvas.bind("<ButtonRelease-1>", self.mouse_release)
self.parent.bind("<Key-Left>", self.open_previous)
self.parent.bind("<Key-Right>", self.open_next)
self.parent.bind("Escape", self.cancel_bbox)
# Labels and Bounding Box Lists Panel
self.listPanel = Frame(self.frame)
self.listPanel.grid(row=0, column=2, sticky=W + N)
self.listBoxNameLabel = Label(self.listPanel,
text="List of Objects").pack(fill=X,
side=TOP)
self.objectListBox = Listbox(self.listPanel, width=40)
self.objectListBox.pack(fill=X, side=TOP)
self.delObjectBtn = Button(self.listPanel,
text="Delete",
command=self.del_bbox)
self.delObjectBtn.pack(fill=X, side=TOP)
self.clearAllBtn = Button(self.listPanel,
text="Clear All",
command=self.clear_bbox)
self.clearAllBtn.pack(fill=X, side=TOP)
self.classesNameLabel = Label(self.listPanel,
text="Classes").pack(fill=X, side=TOP)
self.textBox = Entry(self.listPanel, text="Enter label")
self.textBox.pack(fill=X, side=TOP)
self.addLabelBtn = Button(self.listPanel,
text="+",
command=self.add_label).pack(fill=X,
side=TOP)
self.delLabelBtn = Button(self.listPanel,
text="-",
command=self.del_label).pack(fill=X,
side=TOP)
self.labelListBox = Listbox(self.listPanel)
self.labelListBox.pack(fill=X, side=TOP)
for name in self.names + ['face']:
self.labelListBox.insert(END, str(name))
self.cocoLabels = config.labels_to_names.values()
self.cocoIntVars = []
for idxcoco, label_coco in enumerate(self.cocoLabels):
self.cocoIntVars.append(IntVar())
self.mb.menu.add_checkbutton(label=label_coco,
variable=self.cocoIntVars[idxcoco])
# print(self.cocoIntVars)
# STATUS BAR
self.statusBar = Frame(self.frame, width=500)
self.statusBar.grid(row=1, column=1, sticky=W + N)
self.processingLabel = Label(self.statusBar,
text=" ")
self.processingLabel.pack(side="left", fill=X)
self.imageIdxLabel = Label(self.statusBar,
text=" ")
self.imageIdxLabel.pack(side="right", fill=X)
def open_image(self):
self.filename = filedialog.askopenfilename(
title="Select Image",
filetypes=(("jpeg files", "*.jpg"), ("all files", "*.*")))
# update annotation file
self.anno_filename = '{}_anno.csv'.format(
self.filename.split('/')[-1].split('.')[0])
self.annotation_file = open('annotations/' + self.anno_filename, 'w+')
self.annotation_file.write("")
self.annotation_file.close()
if not self.filename:
return None
self.filenameBuffer = self.filename
self.load_image(self.filenameBuffer)
def open_image_dir(self):
self.imageDir = filedialog.askdirectory(
title="Select Dataset Directory")
if not self.imageDir:
return None
# filename
filename = self.imageDir.split('/')
filename = filename[-2] + '_' + filename[-1]
# update annotation file
self.anno_filename = '{}_anno.csv'.format(filename)
# check if there is already an annotation file
exist = False
all_files = glob('annotations/*.csv')
for file in all_files:
if self.anno_filename in file:
exist = True
break
if exist:
self.annotation_file = open('annotations/' + self.anno_filename,
'r')
# add to self.annoList
for line in self.annotation_file:
line = line.strip()
filename = line.split(',')[0].split('/')[-1]
self.annoList[filename] = line + '\n'
else:
self.annotation_file = open('annotations/' + self.anno_filename,
'w+')
self.annotation_file.write("")
self.annotation_file.close()
self.imageList = os.listdir(self.imageDir)
self.imageList = sorted(self.imageList)
self.imageTotal = len(self.imageList)
self.filename = None
self.imageDirPathBuffer = self.imageDir
self.load_image(self.imageDirPathBuffer + '/' +
self.imageList[self.cur])
def load_image(self, file):
self.img = Image.open(file)
self.imageCur = self.cur + 1
self.imageIdxLabel.config(text=' || Image Number: %d / %d' %
(self.imageCur, self.imageTotal))
# Resize to Pascal VOC format
w, h = self.img.size
if w >= h:
baseW = self.img_size
wpercent = (baseW / float(w))
hsize = int((float(h) * float(wpercent)))
self.img = self.img.resize((baseW, hsize), Image.BICUBIC)
else:
baseH = self.img_size
wpercent = (baseH / float(h))
wsize = int((float(w) * float(wpercent)))
self.img = self.img.resize((wsize, baseH), Image.BICUBIC)
self.tkimg = ImageTk.PhotoImage(self.img)
self.canvas.create_image(0, 0, image=self.tkimg, anchor=NW)
self.clear_bbox()
anno_name = file.split('/')[-1]
if anno_name in self.annoList:
self.load_anno(anno_name)
else:
self.automate()
def open_next(self, event=None):
self.save()
if self.cur < len(self.imageList):
self.cur += 1
self.load_image(self.imageDirPathBuffer + '/' +
self.imageList[self.cur])
self.processingLabel.config(text=" ")
self.processingLabel.update_idletasks()
def open_previous(self, event=None):
self.save()
if self.cur > 0:
self.cur -= 1
self.load_image(self.imageDirPathBuffer + '/' +
self.imageList[self.cur])
self.processingLabel.config(text=" ")
self.processingLabel.update_idletasks()
def save(self):
if self.filenameBuffer is None:
# writer = csv.writer(self.annotation_file)
# key = self.imageDirPathBuffer + '/' + self.imageList[self.cur] + ','
# value = dict()
save_string = self.imageDirPathBuffer + '/' + self.imageList[
self.cur] + ','
for idx, item in enumerate(self.bboxList):
save_string += ','.join(map(
str, self.bboxList[idx])) + ',' + str(
self.objectLabelList[idx] + ',')
# value[str(self.objectLabelList[idx])] = self.bboxList[idx]
else:
save_string = self.imageDirPathBuffer + '/' + self.imageList[
self.cur] + ','
for idx, item in enumerate(self.bboxList):
save_string += ','.join(map(
str, self.bboxList[idx])) + ',' + str(
self.objectLabelList[idx] + ',')
if self.imageList[self.cur] in self.annoList:
if save_string != self.annoList[self.imageList[self.cur]]:
self.annoList[self.imageList[self.cur]] = save_string
filename = save_string.split(',')[0]
# replace
with open('annotations/' + self.anno_filename, 'r') as f:
all_data = f.readlines()
with open('annotations/' + self.anno_filename, 'w') as f:
for i, line in enumerate(all_data):
if filename in line:
f.writelines(save_string + '\n')
else:
f.writelines(line)
else:
self.annotation_file = open('annotations/' + self.anno_filename,
'a')
self.annotation_file.write(save_string + '\n')
self.annotation_file.close()
self.annoList[self.imageList[self.cur]] = save_string
def mouse_click(self, event):
# Check if Updating BBox
if self.canvas.find_enclosed(event.x - 5, event.y - 5, event.x + 5,
event.y + 5):
self.EDIT = True
self.editPointId = int(
self.canvas.find_enclosed(event.x - 5, event.y - 5,
event.x + 5, event.y + 5)[0])
else:
self.EDIT = False
# Set the initial point
if self.EDIT:
idx = self.bboxPointList.index(self.editPointId)
self.editbboxId = self.bboxIdList[math.floor(idx / 4.0)]
self.bboxId = self.editbboxId
pidx = self.bboxIdList.index(self.editbboxId)
pidx = pidx * 4
self.o1 = self.bboxPointList[pidx]
self.o2 = self.bboxPointList[pidx + 1]
self.o3 = self.bboxPointList[pidx + 2]
self.o4 = self.bboxPointList[pidx + 3]
if self.editPointId == self.o1:
a, b, c, d = self.canvas.coords(self.o3)
elif self.editPointId == self.o2:
a, b, c, d = self.canvas.coords(self.o4)
elif self.editPointId == self.o3:
a, b, c, d = self.canvas.coords(self.o1)
elif self.editPointId == self.o4:
a, b, c, d = self.canvas.coords(self.o2)
self.STATE['x'], self.STATE['y'] = int((a + c) / 2), int(
(b + d) / 2)
else:
self.STATE['x'], self.STATE['y'] = event.x, event.y
def mouse_drag(self, event):
self.mouse_move(event)
if self.bboxId:
self.currBboxColor = self.canvas.itemcget(self.bboxId, "outline")
self.canvas.delete(self.bboxId)
self.canvas.delete(self.o1)
self.canvas.delete(self.o2)
self.canvas.delete(self.o3)
self.canvas.delete(self.o4)
if self.EDIT:
self.bboxId = self.canvas.create_rectangle(
self.STATE['x'],
self.STATE['y'],
event.x,
event.y,
width=2,
outline=self.currBboxColor)
else:
self.currBboxColor = config.COLORS[len(self.bboxList) %
len(config.COLORS)]
self.bboxId = self.canvas.create_rectangle(
self.STATE['x'],
self.STATE['y'],
event.x,
event.y,
width=2,
outline=self.currBboxColor)
def mouse_move(self, event):
self.disp.config(text='x: %d, y: %d' % (event.x, event.y))
self.zoom_view(event)
if self.tkimg:
# Horizontal and Vertical Line for precision
if self.hl:
self.canvas.delete(self.hl)
self.hl = self.canvas.create_line(0,
event.y,
self.tkimg.width(),
event.y,
width=2)
if self.vl:
self.canvas.delete(self.vl)
self.vl = self.canvas.create_line(event.x,
0,
event.x,
self.tkimg.height(),
width=2)
# elif (event.x, event.y) in self.bboxBRPointList:
# pass
def mouse_release(self, event):
try:
labelidx = self.labelListBox.curselection()
self.currLabel = self.labelListBox.get(labelidx)
except:
pass
if self.EDIT:
self.update_bbox()
self.EDIT = False
x1, x2 = min(self.STATE['x'], event.x), max(self.STATE['x'], event.x)
y1, y2 = min(self.STATE['y'], event.y), max(self.STATE['y'], event.y)
self.bboxList.append((x1, y1, x2, y2))
o1 = self.canvas.create_oval(x1 - 3,
y1 - 3,
x1 + 3,
y1 + 3,
fill="red")
o2 = self.canvas.create_oval(x2 - 3,
y1 - 3,
x2 + 3,
y1 + 3,
fill="red")
o3 = self.canvas.create_oval(x2 - 3,
y2 - 3,
x2 + 3,
y2 + 3,
fill="red")
o4 = self.canvas.create_oval(x1 - 3,
y2 - 3,
x1 + 3,
y2 + 3,
fill="red")
self.bboxPointList.append(o1)
self.bboxPointList.append(o2)
self.bboxPointList.append(o3)
self.bboxPointList.append(o4)
self.bboxIdList.append(self.bboxId)
self.bboxId = None
self.objectLabelList.append(str(self.currLabel))
self.objectListBox.insert(
END, '(%d, %d) -> (%d, %d)' % (x1, y1, x2, y2) + ': ' +
str(self.currLabel))
self.objectListBox.itemconfig(len(self.bboxIdList) - 1,
fg=self.currBboxColor)
self.currLabel = None
def zoom_view(self, event):
try:
if self.zoomImgId:
self.zoomcanvas.delete(self.zoomImgId)
self.zoomImg = self.img.copy()
self.zoomImgCrop = self.zoomImg.crop(
((event.x - 25), (event.y - 25), (event.x + 25),
(event.y + 25)))
self.zoomImgCrop = self.zoomImgCrop.resize((150, 150))
self.tkZoomImg = ImageTk.PhotoImage(self.zoomImgCrop)
self.zoomImgId = self.zoomcanvas.create_image(0,
0,
image=self.tkZoomImg,
anchor=NW)
hl = self.zoomcanvas.create_line(0, 75, 150, 75, width=2)
vl = self.zoomcanvas.create_line(75, 0, 75, 150, width=2)
except:
pass
def update_bbox(self):
idx = self.bboxIdList.index(self.editbboxId)
self.bboxIdList.pop(idx)
self.bboxList.pop(idx)
self.objectListBox.delete(idx)
self.currLabel = self.objectLabelList[idx]
self.objectLabelList.pop(idx)
idx = idx * 4
self.canvas.delete(self.bboxPointList[idx])
self.canvas.delete(self.bboxPointList[idx + 1])
self.canvas.delete(self.bboxPointList[idx + 2])
self.canvas.delete(self.bboxPointList[idx + 3])
self.bboxPointList.pop(idx)
self.bboxPointList.pop(idx)
self.bboxPointList.pop(idx)
self.bboxPointList.pop(idx)
def cancel_bbox(self, event):
if self.STATE['click'] == 1:
if self.bboxId:
self.canvas.delete(self.bboxId)
self.bboxId = None
self.STATE['click'] = 0
def del_bbox(self):
sel = self.objectListBox.curselection()
if len(sel) != 1:
return
idx = int(sel[0])
self.canvas.delete(self.bboxIdList[idx])
self.canvas.delete(self.bboxPointList[idx * 4])
self.canvas.delete(self.bboxPointList[(idx * 4) + 1])
self.canvas.delete(self.bboxPointList[(idx * 4) + 2])
self.canvas.delete(self.bboxPointList[(idx * 4) + 3])
self.bboxPointList.pop(idx * 4)
self.bboxPointList.pop(idx * 4)
self.bboxPointList.pop(idx * 4)
self.bboxPointList.pop(idx * 4)
self.bboxIdList.pop(idx)
self.bboxList.pop(idx)
self.objectLabelList.pop(idx)
self.objectListBox.delete(idx)
def clear_bbox(self):
for idx in range(len(self.bboxIdList)):
self.canvas.delete(self.bboxIdList[idx])
for idx in range(len(self.bboxPointList)):
self.canvas.delete(self.bboxPointList[idx])
self.objectListBox.delete(0, len(self.bboxList))
self.bboxIdList = []
self.bboxList = []
self.objectLabelList = []
self.bboxPointList = []
def add_label(self):
if self.textBox.get() is not '':
curr_label_list = self.labelListBox.get(0, END)
curr_label_list = list(curr_label_list)
if self.textBox.get() not in curr_label_list:
self.labelListBox.insert(END, str(self.textBox.get()))
self.textBox.delete(0, 'end')
def del_label(self):
labelidx = self.labelListBox.curselection()
self.labelListBox.delete(labelidx)
def add_labels_coco(self):
for listidxcoco, list_label_coco in enumerate(self.cocoLabels):
if self.cocoIntVars[listidxcoco].get():
curr_label_list = self.labelListBox.get(0, END)
curr_label_list = list(curr_label_list)
if list_label_coco not in curr_label_list:
self.labelListBox.insert(END, str(list_label_coco))
def automate(self):
self.clear_bbox()
self.processingLabel.config(text="Processing ")
self.processingLabel.update_idletasks()
open_cv_image0 = np.array(self.img)
# Padded resize
open_cv_image = letterbox(open_cv_image0, new_shape=self.img_size)[0]
# Convert
open_cv_image = open_cv_image[:, :, ::-1].transpose(
2, 0, 1) # BGR to RGB, to 3x416x416
open_cv_image = np.ascontiguousarray(open_cv_image)
img = torch.from_numpy(open_cv_image).to(self.device)
img = img.half() if self.half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# inference object detection
pred = self.object_model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
self.opt.conf_thres,
self.opt.iou_thres,
classes=self.opt.classes,
agnostic=self.opt.agnostic_nms)
# process detections
det = pred[0]
# gn = torch.tensor(open_cv_image0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
open_cv_image0.shape).round()
for idx, (*xyxy, conf, cls) in enumerate(reversed(det)):
if conf < 0.5:
continue
xyxy = torch.tensor(xyxy).view(-1).cpu().numpy().astype(np.int)
label = config.labels_to_names[int(cls)]
curr_label_list = self.labelListBox.get(0, END)
curr_label_list = list(curr_label_list)
if label not in curr_label_list:
continue
b = xyxy.tolist()
self.bboxId = self.canvas.create_rectangle(
b[0],
b[1],
b[2],
b[3],
width=2,
outline=config.COLORS[len(self.bboxList) %
len(config.COLORS)])
self.bboxList.append((b[0], b[1], b[2], b[3]))
o1 = self.canvas.create_oval(b[0] - 3,
b[1] - 3,
b[0] + 3,
b[1] + 3,
fill="red")
o2 = self.canvas.create_oval(b[2] - 3,
b[1] - 3,
b[2] + 3,
b[1] + 3,
fill="red")
o3 = self.canvas.create_oval(b[2] - 3,
b[3] - 3,
b[2] + 3,
b[3] + 3,
fill="red")
o4 = self.canvas.create_oval(b[0] - 3,
b[3] - 3,
b[0] + 3,
b[3] + 3,
fill="red")
self.bboxPointList.append(o1)
self.bboxPointList.append(o2)
self.bboxPointList.append(o3)
self.bboxPointList.append(o4)
self.bboxIdList.append(self.bboxId)
self.bboxId = None
self.objectLabelList.append(str(label))
self.objectListBox.insert(
END, '(%d, %d) -> (%d, %d)' % (b[0], b[1], b[2], b[3]) +
': ' + str(label))
self.objectListBox.itemconfig(
len(self.bboxIdList) - 1,
fg=config.COLORS[(len(self.bboxIdList) - 1) %
len(config.COLORS)])
# inference faces
frame = Image.fromarray(open_cv_image0[:, :, ::-1])
# detect faces
boxes, probs = self.face_model.detect(frame, landmarks=False)
if boxes is None:
boxes = []
probs = []
for box, conf in zip(boxes, probs):
if conf < 0.5:
continue
# box in xyxy format
b = box.astype(np.int).tolist()
label = 'face'
curr_label_list = self.labelListBox.get(0, END)
curr_label_list = list(curr_label_list)
if label not in curr_label_list:
continue
self.bboxId = self.canvas.create_rectangle(
b[0],
b[1],
b[2],
b[3],
width=2,
outline=config.COLORS[len(self.bboxList) % len(config.COLORS)])
self.bboxList.append((b[0], b[1], b[2], b[3]))
o1 = self.canvas.create_oval(b[0] - 3,
b[1] - 3,
b[0] + 3,
b[1] + 3,
fill="red")
o2 = self.canvas.create_oval(b[2] - 3,
b[1] - 3,
b[2] + 3,
b[1] + 3,
fill="red")
o3 = self.canvas.create_oval(b[2] - 3,
b[3] - 3,
b[2] + 3,
b[3] + 3,
fill="red")
o4 = self.canvas.create_oval(b[0] - 3,
b[3] - 3,
b[0] + 3,
b[3] + 3,
fill="red")
self.bboxPointList.append(o1)
self.bboxPointList.append(o2)
self.bboxPointList.append(o3)
self.bboxPointList.append(o4)
self.bboxIdList.append(self.bboxId)
self.bboxId = None
self.objectLabelList.append(str(label))
self.objectListBox.insert(
END, '(%d, %d) -> (%d, %d)' % (b[0], b[1], b[2], b[3]) + ': ' +
str(label))
self.objectListBox.itemconfig(
len(self.bboxIdList) - 1,
fg=config.COLORS[(len(self.bboxIdList) - 1) %
len(config.COLORS)])
self.processingLabel.config(text="Done")
def load_anno(self, file):
annotations = self.annoList[file].strip().split(',')
for i in range(1, len(annotations), 5):
if annotations[i] == '' or annotations[i] == '\n':
continue
x1, y1, x2, y2 = annotations[i:i + 4]
label = annotations[i + 4]
b = [int(x1), int(y1), int(x2), int(y2)]
self.bboxId = self.canvas.create_rectangle(
b[0],
b[1],
b[2],
b[3],
width=2,
outline=config.COLORS[len(self.bboxList) % len(config.COLORS)])
self.bboxList.append((b[0], b[1], b[2], b[3]))
o1 = self.canvas.create_oval(b[0] - 3,
b[1] - 3,
b[0] + 3,
b[1] + 3,
fill="red")
o2 = self.canvas.create_oval(b[2] - 3,
b[1] - 3,
b[2] + 3,
b[1] + 3,
fill="red")
o3 = self.canvas.create_oval(b[2] - 3,
b[3] - 3,
b[2] + 3,
b[3] + 3,
fill="red")
o4 = self.canvas.create_oval(b[0] - 3,
b[3] - 3,
b[0] + 3,
b[3] + 3,
fill="red")
self.bboxPointList.append(o1)
self.bboxPointList.append(o2)
self.bboxPointList.append(o3)
self.bboxPointList.append(o4)
self.bboxIdList.append(self.bboxId)
self.bboxId = None
self.objectLabelList.append(str(label))
self.objectListBox.insert(
END, '(%d, %d) -> (%d, %d)' % (b[0], b[1], b[2], b[3]) + ': ' +
str(label))
self.objectListBox.itemconfig(
len(self.bboxIdList) - 1,
fg=config.COLORS[(len(self.bboxIdList) - 1) %
len(config.COLORS)])
extract_face(
img,
box,
save_path=save_path.joinpath(f'{img_name}_detected_face_{i}.jpg'))
img_draw.save(save_path.joinpath(f'{img_name}_annotate_faces.jpg'))
return True
def draw_img(model, img_dir, save_dir):
images = []
for suffix in ['jpg', 'jpeg', 'png']:
images.extend(glob.glob(str(img_dir.joinpath(f'*/*.{suffix}'))))
for img in images:
img_path = pathlib.Path(img)
save_path = save_dir.joinpath(img_path.parent.stem)
extract_img(model, img_path, save_path)
if __name__ == "__main__":
model = MTCNN(keep_all=True)
resnet = InceptionResnetV1(pretrained='vggface2')
data_dir = pathlib.Path('dataset/images/')
test_dir = pathlib.Path('dataset/test/')
draw_img(model, data_dir, test_dir)
def collate_fn(x):
return x[0]
# dataset = datasets.ImageFolder(data_dir)
# loader = DataLoader(dataset, collate_fn=collate_fn)
class faceapp():
def __init__(self, size, margin, weights, device, select_largest, images):
print("Loading network.....")
current = time.time()
self.device = device
self.mtcnn = MTCNN(image_size=size,
margin=margin,
keep_all=False,
device=device,
select_largest=True)
self.model = InceptionResnetV1(pretrained='vggface2', device=device)
state_dict = torch.load(weights, map_location="cpu")
self.model.load_state_dict(state_dict)
self.images = images
elapse = time.time() - current
print("Network successfully loaded, process time: %.2fs" % elapse)
# eval mode
self.model.eval()
def realtime(self):
cap = cv2.VideoCapture(0)
cap.set(3, 1280)
cap.set(4, 720)
prev_frame_time = 0
new_frame_time = 0
assert cap.isOpened(), 'Cannot capture source'
frames = 0
start = time.time()
while cap.isOpened():
ret, frame = cap.read()
if ret:
boxes, probs, points = self.mtcnn.detect(frame, landmarks=True)
if boxes is None:
continue
dst = draw_boxes(frame, boxes, probs)
font = cv2.FONT_HERSHEY_PLAIN
new_frame_time = time.time()
try:
fps = 1 / (new_frame_time - prev_frame_time)
except ZeroDivisionError:
continue
prev_frame_time = new_frame_time
fps = int(fps)
fps = str(fps)
cv2.putText(dst, "FPS:%s" % fps, (10, 25), font, 1,
[0, 0, 255], 1)
cv2.imshow('temp', dst)
key = cv2.waitKey(1)
if key & 0xFF == ord('q'):
break
cv2.destroyAllWindows()
def collect(self, img):
img_cropped, prob = self.mtcnn(img, return_prob=True)
if img_cropped is not None:
boxes, probs, points = self.mtcnn.detect(img, landmarks=True)
# draw bbox in image
dst = draw_boxes(img, boxes, probs)
img_cropped = img_cropped.to(self.device)
# Calculate embedding (unsqueeze to add batch dimension)
img_embedding = self.model(img_cropped.unsqueeze(0)).to('cpu')
return ErrEnum.No_problem, img_embedding, dst
else:
print("No face detected, Try again!!!")
return ErrEnum.No_face_detected, None, None
def recognize(self, img, threshold):
img_cropped, prob = self.mtcnn(img, return_prob=True)
if img_cropped is not None:
boxes, probs, points = self.mtcnn.detect(img, landmarks=True)
img_cropped = img_cropped.to(self.device)
# Calculate embedding (unsqueeze to add batch dimension)
img_embedding = self.model(img_cropped.unsqueeze(0)).to('cpu')
smallest_dist = 999
embs = glob.glob(osp.realpath(self.images) + '\\*.embs')
for emb in embs:
dist = calc_dist(img_embedding.detach().numpy(),
emb2numpy(emb))
if smallest_dist > dist:
smallest_dist = dist
recog_name = emb.split('\\')[-1].split('.')[0]
if smallest_dist > threshold:
# print('recognize unreliable!')
return ErrEnum.Unreliable_recognize, "", None
else:
print("#" * 50)
print('Staff identified! Welcome: {}'.format(recog_name))
print("login time is: {}".format(
datetime.datetime.now().strftime('%Y/%m/%d %H:%M:%S')))
print("#" * 50)
dst = draw_boxes(img, boxes, recog_name)
cv2.putText(dst, str(smallest_dist), (10, 710),
cv2.FONT_HERSHEY_PLAIN, 1, [225, 255, 255], 1)
return ErrEnum.No_problem, recog_name, dst
else:
print("No face detected, Try again!!!")
return ErrEnum.No_face_detected, None, None