- 主界面
class Display():
def __init__(self):
self._init_root(WIDE=800, HEIGHT=500)
self.train_args=None
self.train_paras_info=""
#flag for end train
self.train_end=False
#video stream source type option,"camera" or "video"
self.video_type=""
#video file path
self.video_path=''
#use to exit video paly
self.show_state=False
#play speed
self.speed_level=0
def run(self):
self.root.mainloop()
def _init_root(self, WIDE=600, HEIGHT=400):
self.WIDE,self.HEIGHT=WIDE,HEIGHT
#creat window
self.root=tk.Tk()
self.root.title("Single Shot MultiBox Detector")
self.root.config(background="white")
self.root.geometry(str(self.WIDE) + "x" + str(self.HEIGHT)+"-500-200")
self.root.resizable(0, 0)
#self.root.iconbitmap("./image/icon.ico")
self.root.attributes("-alpha", 0.95)
#self.root.protocol("WM_DELETE_WINDOW", call) #when window close ,exec call
# creat video area
self._crerat_video_area()
#state area
self.text = tk.Label(self.root, text="train or detect 1.setting 2.start 3.end")
self.text.place(x=1, y=self.video_h-32, w=self.video_w, h=30)
#detailed information area
self.detial_info=tk.Label(self.root,text="welcome",bg="white",fg="green",justify="left")
self.detial_info.place(x=3 / 4 * (self.WIDE) + 5, y= 2, w=1/4*self.WIDE-8, h=1/2*self.HEIGHT)
# creat button
button_w,button_h=100,30
button_x ,button_y= 7 / 8 * self.WIDE - 1 / 2 * button_w,6/11 * self.HEIGHT
train_setting=tk.Button(self.root, text="train setting", relief="groove", command=self._train_setting)
train_setting.place(x=button_x, y= button_y, w=button_w, h=button_h)
train_start = tk.Button(self.root, text="start", relief="groove", command=self._train_start)
train_start.place(x=button_x, y= button_y+button_h, w=1/2*button_w, h=button_h)
train_end= tk.Button(self.root, text="end ", relief="groove", command=self._train_end, )
train_end.place(x=button_x+1/2*button_w, y= button_y+button_h, w=1/2*button_w, h=button_h)
detect_setting = tk.Button(self.root, text="detect setting", relief="groove", command=self._detect_setting, )
detect_setting.place(x=button_x, y= button_y+100, w=button_w, h=button_h)
detect_start = tk.Button(self.root, text="start", relief="groove", command=self._detect_start)
detect_start.place(x=button_x, y= button_y+100+button_h, w=1/2*button_w, h=button_h)
detect_end= tk.Button(self.root, text="end ", relief="groove", command=self._detect_end, )
detect_end.place(x=button_x+1/2*button_w, y= button_y+100+button_h, w=1/2*button_w, h=button_h)
detect_slow=tk.Button(self.root, text="slow -", relief="groove", command=self._detect_slow,)
detect_slow.place(x=button_x, y= button_y+100+2*button_h, w=1/2*button_w, h=button_h)
detect_quick=tk.Button(self.root, text="quick +", relief="groove", command=self._detect_quick,)
detect_quick.place(x=button_x+1/2*button_w, y= button_y+100+2*button_h, w=1/2*button_w, h=button_h)
#creat video pr text area.when swith usage,destroy and creat new
def _crerat_video_area(self):
self.video_w, self.video_h= int(3 / 4 * self.WIDE) - 1, int(self.HEIGHT) - 2
self.video = tk.Label(self.root, bg="black", fg="green",cursor="cross")
self.video.place(x=1, y=1, w=self.video_w, h=self.video_h-35)
"""
@ train methods ################################################
"""
def _train_setting(self):
self.video.destroy()
self._crerat_video_area()
self.train_args,self.train_paras_info=TrainSetting().run()
self.video.config(image=None,text=self.train_paras_info,justify="left")#,wraplength=500
self.video.update()
def _train_end(self):
self.train_end=True
def _train_start(self):
import train
train.train(self.train_args)
"""
@ detect methods ################################################
"""
def _detect_setting(self):
self.video_type,self.video_path=DetectSetting().run()
def _detect_quick(self):
self.speed_level+=1
def _detect_slow(self):
if self.speed_level>0:
self.speed_level-=1
def _detect_end(self):
self.show_state=False
def _detect_start(self):
self.video.destroy()
self._crerat_video_area()
self.show_state = True
if self.video_type=="camera":
cap=cv2.VideoCapture(0,cv2.CAP_DSHOW)
elif self.video_type=="video":
cap = cv2.VideoCapture(self.video_path)
else:
self.text.config(text="No video stream selected")
self.text.update()
return
# mb1-ssd models/mobilenet-v1-ssd-mp-0_675.pth models/voc-model-labels.txt
# or mb2-ssd-lite models/mb2-ssd-lite-mp-0_686.pth models/voc-model-labels.txt
net_type = "mb2-ssd-lite"
model_path = r"models\mb2-ssd-lite-mp-0_686.pth"
label_path = r"models\voc-model-labels.txt "
self.class_names = [name.strip() for name in open(label_path).readlines()]
num_classes = len(self.class_names)
if net_type == 'mb1-ssd':
net = create_mobilenetv1_ssd(len(self.class_names), is_test=True)
net.load(model_path)
self.predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(self.class_names), is_test=True)
net.load(model_path)
self.predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200)
ret, orig_image = cap.read()
w,h=orig_image.shape[1],orig_image.shape[0]
h=int(h*(self.video_w/w))
w=self.video_w
counter=0
while self.show_state:
ret, orig_image = cap.read()
if orig_image is None:
continue
#adjust play speed
if counter!=self.speed_level:
counter+=1
continue
else:
counter=0
orig_image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
ts=time.time()
boxes, labels, probs = self.predictor.predict(orig_image, 10, 0.4)
interval = time.time()-ts
show_info=""
for i in range(boxes.size(0)):
box = boxes[i, :]
label = f"{self.class_names[labels[i]]}: {probs[i]:.2f}"
cv2.rectangle(orig_image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (255, 255, 0), 4)
cv2.putText(orig_image, label,
(int(box[0]) + 20, int(box[1]) + 40),
cv2.FONT_HERSHEY_SIMPLEX,
1, # font scale
(255, 0, 255),
2) # line type
show_info+=label+"\n"+str(int(box[0]))+"," +str(int(box[1]))+","+str(int(box[2]))+","+ str(int(box[3]))+"\n"
#convert the image format to ImageTk ,use to show in tk label
current_image = Image.fromarray(orig_image,"RGB").resize((w,h),Image.ANTIALIAS)
img= ImageTk.PhotoImage(image=current_image)
self.video.config(image=img)
self.text.config(text='Time: {:.2f}s, Detect Objects: {:d}.'.format(interval, labels.size(0)))
self.detial_info.config(text=show_info)
self.video.update()
self.text.update()
self.detial_info.update()
#cv2.waitKey(100)
if ret is False:
break
cap.release()
- 训练参数设置
#train dialog setting window
class TrainSetting():
def __init__(self,WIN_W=600,WIN_H=430):
# use argparse lib creat namespace ,to save train args
self.args = argparse.ArgumentParser().parse_args()
self.args.labels_file=r".\models\voc-model-labels.txt"
self.args.image_id_file=os.path.join(r"..\pascalvoc\VOCdevkit\VOC2012","ImageSets","Main","trainval.txt")
self.args.image_folder=os.path.join("..\pascalvoc\VOCdevkit\VOC2012","JPEGImages")
self.args.annotations_folder=os.path.join("..\pascalvoc\VOCdevkit\VOC2012","Annotations")
# self.args.datasets ="..\pascalvoc\VOCdevkit\VOC2012"
# self.args.validation_dataset = "..\pascalvoc\VOCdevkit\VOC2007"
self.args.checkpoint_folder='models/'
self.args.batch_size = 24
self.args.num_epochs = 200
self.args.num_workers = 4
self.args.use_cuda = True
self.args.net = 'mb3-small-ssd-lite'
self.args.base_net = "" # os.path.join("models", "mobilenet_v1_with_relu_69_5.pth")
self.args.scheduler = "cosine"
self.args.lr = 0.01
self.args.t_max = 200 # T_max value for Cosine Annealing Scheduler
#self.args.balance_data='store_true'
self.args.freeze_base_net="store_true"
self.args.freeze_net="store_true"
#self.args.mb2_width_mult=1.0
self.args.momentum=0.9
self.args.weight_decay=5e-4
self.args.gamma=0.1
self.args.base_net_lr=None
self.args.extra_layers_lr=None
self.args.pretrained_ssd=None
self.args.resume=None
self.args.scheduler="multi-step"
self.args.milestones="80,100"
# self.args.validation_epochs=5
win_w, win_h = WIN_W, WIN_H
self.dialog = tk.Tk()
self.dialog.title("train setting")
self.dialog.geometry(str(win_w) + "x" + str(win_h))
self.dialog.attributes("-alpha", 0.95)
#button action
def _choose_labels_file():
self.args.labels_file = tkinter.filedialog.askopenfilename()
t_labels_file.insert(0,self.args.labels_file)
def _choose_image_id_file():
self.args.image_id_file= tkinter.filedialog.askopenfilename()
t_image_id_file.insert(0,self.args.image_id_file)
def _choose_image_folder():
self.args.image_folder=tkinter.filedialog.askopenfilename()
t_image_folder.insert(0,self.args.image_folder)
def _choose_annotation_folder():
self.args.annotations_folder=tkinter.filedialog.askopenfilename()
t_annotation_folder.insert(0,self.args.annotations_folder)
def _choose_checkpoint_folder():
self.args.checkpoint_folder=tkinter.filedialog.askopenfilename()
t_checkpoint_folder.insert(0,self.args.checkpoint_folder)
def _choose_cpu():
self.args.use_cuda=False
def _choos_gpu():
self.args.use_cuda=True
def _exit_dialog():
self.args.batch_size=int(t_batch_size.get())
self.args.num_epochs=int(t_num_epochs.get())
self.args.num_workers=int(t_num_workers.get())
#qiut to return args,then destroy window
self.dialog.quit()
# label
tk.Label(self.dialog, text="labels file").place(x=5, y=10, w=100, h=30)
tk.Label(self.dialog, text="image id file").place(x=10, y=50, w=100, h=30)
tk.Label(self.dialog, text="image folder").place(x=3, y=90, w=120, h=30)
tk.Label(self.dialog, text="annotation folder").place(x=0, y=130, w=150, h=30)
tk.Label(self.dialog, text="batch size").place(x=5, y=170, w=100, h=30)
tk.Label(self.dialog, text="num_epochs").place(x=10, y=210, w=100, h=30)
tk.Label(self.dialog, text="num_workers").place(x=10, y=250, w=100, h=30)
tk.Label(self.dialog, text="checkpoint_folder").place(x=0, y=290, w=150, h=30)
tk.Label(self.dialog, text="mode").place(x=10, y=330, w=100, h=30)
#Entry
t_labels_file=tk.Entry(self.dialog)
t_labels_file.insert(0,self.args.labels_file)
t_labels_file.place(x=140, y=10, w=300, h=30)
t_image_id_file=tk.Entry(self.dialog)
t_image_id_file.insert(0,self.args.image_id_file)
t_image_id_file.place(x=140, y=50, w=300, h=30)
t_image_folder=tk.Entry(self.dialog)
t_image_folder.insert(0,self.args.image_folder)
t_image_folder.place(x=140, y=90, w=300, h=30)
t_annotation_folder=tk.Entry(self.dialog)
t_annotation_folder.insert(0,self.args.annotations_folder)
t_annotation_folder.place(x=140,y=130,w=300,h=30)
t_batch_size=tk.Entry(self.dialog)
t_batch_size.insert(0,self.args.batch_size)
t_batch_size.place(x=140,y=170,w=300,h=30)
t_num_epochs=tk.Entry(self.dialog)
t_num_epochs.insert(0,self.args.num_epochs)
t_num_epochs.place(x=140,y=210,w=300,h=30)
t_num_workers=tk.Entry(self.dialog)
t_num_workers.insert(0,self.args.num_workers)
t_num_workers.place(x=140,y=250,w=300,h=30)
t_checkpoint_folder=tk.Entry(self.dialog)
t_checkpoint_folder.insert(0,self.args.checkpoint_folder)
t_checkpoint_folder.place(x=140,y=290,w=300,h=30)
#button
tk.Button(self.dialog, text="choose file",relief='groove', command=_choose_labels_file)\
.place(x=win_w-130, y=10, w=120, h=30)
tk.Button(self.dialog, text="choose file", relief='groove',command=_choose_image_id_file)\
.place(x=win_w-130, y=50, w=120, h=30)
tk.Button(self.dialog, text="choose folder", relief='groove',command=_choose_image_folder)\
.place(x=win_w-130, y=90, w=120, h=30)
tk.Button(self.dialog, text="choose folder", relief='groove',command=_choose_annotation_folder)\
.place(x=win_w-130, y=130, w=120, h=30)
tk.Button(self.dialog, text="choose folder", relief='groove',command=_choose_checkpoint_folder)\
.place(x=win_w-130, y=290, w=120, h=30)
b_voc=tk.Checkbutton(self.dialog, text="cpu", command=_choose_cpu)
b_voc.select()
b_voc.place(x=140,y=330,w=100,h=30)
tk.Checkbutton(self.dialog, text="gpu", command=_choos_gpu)\
.place(x=250,y=330,w=100,h=30)
tk.Button(self.dialog, text="ok", relief='groove',command=_exit_dialog)\
.place(x=1/2*win_w-40, y=380, w=80, h=30)
def run(self):
self.dialog.mainloop()
self.dialog.destroy()
self._convert_info()
return self.args,self.info
def _convert_info(self):
self.info=""
args_dict=self.args.__dict__
for i,(key,value) in enumerate(args_dict.items()):
self.info+=str(key)+":"+str(value)+"\n"
if i==10:
return
- 检测相关设置
class DetectSetting():
def __init__(self):
self.video_type = ""
self.video_path = ""
win_w,win_h=200,100
self.dialog = tk.Tk()
self.dialog.title("select video stream")
self.dialog.geometry(str(win_w)+"x"+str(win_h))
self.dialog.attributes("-alpha", 0.95)
def _use_camera():
self.video_type="camera"
self.video_path=""
self.dialog.quit()
def _use_video():
self.video_type="video"
self.video_path = tkinter.filedialog.askopenfilename()
self.dialog.quit()
use_camera = tk.Button(self.dialog, text="use camera", command=_use_camera,relief='groove')
use_camera.place(x=10, y=1/2*win_h-15, w=70, h=30)
use_vedio = tk.Button(self.dialog, text="user video", command=_use_video,relief='groove')
use_vedio.place(x=win_w-70-10, y=1/2*win_h-15, w=70, h=30)
def run(self):
self.dialog.mainloop()
self.dialog.destroy()
return self.video_type,self.video_path
细节:
1子窗口传递返回参数
窗口.quit(),执mainloop()之后的return语句返回参数
2.Label只能显示 图像或文字的一种,图像优先级高,从图像切换回文字会报错
解决思路:用函数初始化Label,切换时先destroy(),再在原来位置初始化一个新label
注
1.绑定事件
b=Button(root,bg='green')
b.place(x=9/10*WIDE+1,y=47/60*(HEIGHT),w=10,h=10)
#绑定鼠标左键击键
b.bind('<Button-1>', onclick)
#绑定回车键击键
root.bind('<Return>', onclick)