Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Features added #15

Merged
merged 9 commits into from
Apr 11, 2020
Merged

Features added #15

Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
16ebe0d
features added
QuickLearner171998 Mar 29, 2020
9773ba0
Update main.py
QuickLearner171998 Mar 29, 2020
2f76818
added option to switch between models
QuickLearner171998 Mar 31, 2020
140931f
added separate file for tensorflow labels
QuickLearner171998 Mar 31, 2020
52ed5f4
Create README
QuickLearner171998 Mar 31, 2020
c5d3bd1
Create readme
QuickLearner171998 Mar 31, 2020
273c7f2
Update README.md
QuickLearner171998 Mar 31, 2020
20759ce
minor bugs fixed
QuickLearner171998 Apr 11, 2020
f3ae901
remove print
QuickLearner171998 Apr 11, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 130 additions & 35 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,7 @@
from tkinter import *
from tkinter import filedialog
from PIL import Image, ImageTk

import keras

from keras_retinanet import models
from keras_retinanet.utils.image import preprocess_image

Expand All @@ -20,26 +18,26 @@
import numpy as np
import tensorflow as tf
import config
import tf_config
import math
# for saving to xml files
from pascal_voc_writer import Writer


def get_session():
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
return tf.Session(config=config)


keras.backend.tensorflow_backend.set_session(get_session())

model_path = os.path.join('.', 'snapshots', 'resnet50_coco_best_v2.1.0.h5')

model = models.load_model(model_path, backbone_name='resnet50')
# make sure the file is inside semi-auto-image-annotation-tool-master
import pathlib
cur_path = pathlib.Path(__file__).parent.absolute()
# cur_path = r"semi-auto-image-annotation-tool-master"
sys.path.append(cur_path)
os.chdir(cur_path)
QuickLearner171998 marked this conversation as resolved.
Show resolved Hide resolved


class MainGUI:
def __init__(self, master):

# to choose between keras or tensorflow models
self.kera = 1 # default
self.tensorf = 0
QuickLearner171998 marked this conversation as resolved.
Show resolved Hide resolved
self.models_dir = '' # gets updated as per user choice
self.model_path = ''
self.parent = master
self.parent.title("Semi Automatic Image Annotation Tool")
self.frame = Frame(self.parent)
Expand Down Expand Up @@ -77,9 +75,11 @@ def __init__(self, master):
self.filenameBuffer = None
self.objectLabelList = []
self.EDIT = False
self.autoSuggest = StringVar()
self.writer = None
# default thresh can be modified by the user later
self.thresh = 0.5
self.org_h = 0
self.org_w = 0
# initialize mouse state
self.STATE = {'x': 0, 'y': 0}
self.STATE_COCO = {'click': 0}
Expand All @@ -99,26 +99,43 @@ def __init__(self, master):
self.openBtn.grid(columnspan=2, sticky=W + E)
self.openDirBtn = Button(self.ctrlPanel, text='Open Dir', command=self.open_image_dir)
self.openDirBtn.grid(columnspan=2, sticky = W + E)

self.nextBtn = Button(self.ctrlPanel, text='Next -->', command=self.open_next)
self.nextBtn.grid(columnspan=2, sticky=W + E)
self.previousBtn = Button(self.ctrlPanel, text='<-- Previous', command=self.open_previous)
self.previousBtn.grid(columnspan=2, sticky=W + E)
self.saveBtn = Button(self.ctrlPanel, text='Save', command=self.save)
self.saveBtn.grid(columnspan=2, sticky=W + E)
self.semiAutoBtn = Button(self.ctrlPanel, text="Show Suggestions", command=self.automate)
self.autoManualLabel = Label(self.ctrlPanel, text="Suggestion Mode")
self.autoManualLabel.grid(columnspan=2, sticky=W + E)
self.radioBtnAuto = Radiobutton(self.ctrlPanel, text="Auto", variable=self.autoSuggest, value=1)
self.radioBtnAuto.grid(row=7, column=0, sticky=W + E)
self.radioBtnManual = Radiobutton(self.ctrlPanel, text="Manual", variable=self.autoSuggest, value=2)
self.radioBtnManual.grid(row=7, column=1, sticky=W + E)
self.semiAutoBtn = Button(self.ctrlPanel, text="Detect", command=self.automate)
self.semiAutoBtn.grid(columnspan=2, sticky=W + E)
self.disp = Label(self.ctrlPanel, text='Coordinates:')
self.disp.grid(columnspan=2, sticky=W + E)

self.mb = Menubutton(self.ctrlPanel, text="COCO Classes for Suggestions", relief=RAISED)
self.mb.grid(columnspan=2, sticky=W + E)
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.grid(columnspan=2, sticky=W + E)
# option to add all classes to list
self.addCocoBtnAllClasses = Button(self.ctrlPanel, text="Add All Classes", command=self.add_all_classes)
self.addCocoBtnAllClasses.grid(columnspan=2, sticky=W + E)

# options to add different models
self.mb1 = Menubutton(self.ctrlPanel, text="Select models from here", relief=RAISED)
self.mb1.grid(columnspan=2, sticky=W + E)
self.mb1.menu = Menu(self.mb1, tearoff=0)
self.mb1["menu"] = self.mb1.menu

self.addModelBtn = Button(self.ctrlPanel, text="Add model", command=self.add_model)
self.addModelBtn.grid(columnspan=2, sticky=W + E)

self.zoomPanelLabel = Label(self.ctrlPanel, text="Precision View Panel")
self.zoomPanelLabel.grid(columnspan=2, sticky=W + E)
self.zoomcanvas = Canvas(self.ctrlPanel, width=150, height=150)
Expand Down Expand Up @@ -150,23 +167,30 @@ def __init__(self, master):
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)
# add option to set threshold

self.addThresh = Label(self.listPanel, text="Threshold").pack(fill=X, side=TOP)
self.textBoxTh = Entry(self.listPanel, text="Enter threshold value")
self.textBoxTh.pack(fill=X, side=TOP)
self.enterthresh = Button(self.listPanel, text="Set", command=self.changeThresh).pack(fill=X, side=TOP)

self.labelListBox = Listbox(self.listPanel)
self.labelListBox.pack(fill=X, side=TOP)
if self.kera:
self.cocoLabels = config.labels_to_names.values()
else:
self.cocoLabels = tf_config.labels_to_names.values()

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)

self.modelIntVars = []
for idxmodel, modelname in enumerate(self.available_models()):
self.modelIntVars.append(IntVar())
self.mb1.menu.add_checkbutton(label=modelname, variable=self.modelIntVars[idxmodel])

# STATUS BAR
self.statusBar = Frame(self.frame, width=500)
self.statusBar.grid(row=1, column=1, sticky=W + N)
Expand All @@ -175,15 +199,27 @@ def __init__(self, master):
self.imageIdxLabel = Label(self.statusBar, text=" ")
self.imageIdxLabel.pack(side="right", fill=X)

def get_session(self):
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
return tf.Session(config=config)

def available_models(self):
self.models_dir = os.path.join(cur_path, 'snapshots')
# only for keras and tf
model_categ = os.listdir(self.models_dir)
# creating all model options list
model_names = []
for categ in model_categ:
for name in os.listdir(os.path.join(self.models_dir , categ)):
model_names.append(os.path.join(categ,name))
QuickLearner171998 marked this conversation as resolved.
Show resolved Hide resolved
return model_names


def changeThresh(self):
# to update the threshold
if(float(self.textBoxTh.get()) >0 and float(self.textBoxTh.get()) <1):
self.thresh = float(self.textBoxTh.get())



def open_image(self):
self.filename = filedialog.askopenfilename(title="Select Image", filetypes=(("jpeg files", "*.jpg"),
("all files", "*.*")))
Expand All @@ -203,12 +239,16 @@ def open_image_dir(self):
self.imageDirPathBuffer = self.imageDir
self.load_image(self.imageDirPathBuffer + '/' + self.imageList[self.cur])

def open_video_file(self):
pass

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
self.org_w, self.org_h = self.img.size
if w >= h:
baseW = 500
wpercent = (baseW / float(w))
Expand All @@ -231,6 +271,8 @@ def open_next(self, event=None):
self.load_image(self.imageDirPathBuffer + '/' + self.imageList[self.cur])
self.processingLabel.config(text=" ")
self.processingLabel.update_idletasks()
if self.autoSuggest.get() == str(1):
self.automate()

def open_previous(self, event=None):
self.save()
Expand All @@ -239,6 +281,8 @@ def open_previous(self, event=None):
self.load_image(self.imageDirPathBuffer + '/' + self.imageList[self.cur])
self.processingLabel.config(text=" ")
self.processingLabel.update_idletasks()
if self.autoSuggest.get() == str(1):
self.automate()

def save(self):
if self.filenameBuffer is None:
Expand Down Expand Up @@ -448,6 +492,24 @@ def del_label(self):
labelidx = self.labelListBox.curselection()
self.labelListBox.delete(labelidx)

def add_model(self):
for listidxmodel, list_model_name in enumerate(self.available_models()):
if(self.modelIntVars[listidxmodel].get()):
# check which model is it keras or tensorflow
self.model_path = os.path.join(self.models_dir,list_model_name)
# if its Tensorflow model then modify path
if('keras' in list_model_name):
self.kera = 1
self.tensorf = 0
elif('tensorflow' in list_model_name):
self.model_path = os.path.join(self.model_path,'frozen_inference_graph.pb')
self.kera = 0
self.tensorf = 1
# change cocoLabels corresponding to tensorflow
self.cocoLabels = tf_config.labels_to_names.values()
break


def add_labels_coco(self):
for listidxcoco, list_label_coco in enumerate(self.cocoLabels):
if self.cocoIntVars[listidxcoco].get():
Expand All @@ -458,6 +520,7 @@ def add_labels_coco(self):

def add_all_classes(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:
Expand All @@ -468,24 +531,55 @@ def automate(self):
self.processingLabel.update_idletasks()
open_cv_image = np.array(self.img)
# Convert RGB to BGR
opencvImage= open_cv_image[:, :, ::-1].copy()
# opencvImage = cv2.cvtColor(np.array(self.img), cv2.COLOR_RGB2BGR)
# self.thresh = self.thresh
opencvImage = open_cv_image[:, :, ::-1].copy()
# if tensorflow
if self.tensorf :
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(self.model_path, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')

sess = tf.Session(graph=detection_graph)

image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = detection_graph.get_tensor_by_name('num_detections:0')

image_expanded = np.expand_dims(opencvImage, axis=0)
(boxes, scores, labels, num) = sess.run(
[detection_boxes, detection_scores, detection_classes, num_detections],
feed_dict={image_tensor: image_expanded})
config_labels = tf_config.labels_to_names
m_name = os.path.split((os.path.split(self.model_path)[0]))[1]


image = preprocess_image(opencvImage)
boxes, scores, labels = model.predict_on_batch(np.expand_dims(image, axis=0))
else:
keras.backend.tensorflow_backend.set_session(self.get_session())
model_path = self.model_path
model = models.load_model(model_path, backbone_name='resnet50')
image = preprocess_image(opencvImage)
boxes, scores, labels = model.predict_on_batch(np.expand_dims(image, axis=0))
config_labels = config.labels_to_names
m_name = os.path.split(self.model_path)[1]
for idx, (box, label, score) in enumerate(zip(boxes[0], labels[0], scores[0])):
curr_label_list = self.labelListBox.get(0, END)
curr_label_list = list(curr_label_list)
if score < self.thresh:
continue

if config.labels_to_names[label] not in curr_label_list:
if config_labels[label] not in curr_label_list:
continue

b = box.astype(int)

b = box
# only if using tf models as keras and tensorflow have different coordinate order
if(self.tensorf):
w, h = self.img.size
(b[0],b[1],b[2],b[3]) = (b[1]*w, b[0]*h, b[3]*w, b[2]*h)
b = b.astype(int)
self.bboxId = self.canvas.create_rectangle(b[0], b[1],
b[2], b[3],
width=2,
Expand All @@ -501,9 +595,10 @@ def automate(self):
self.bboxPointList.append(o4)
self.bboxIdList.append(self.bboxId)
self.bboxId = None
self.objectLabelList.append(str(config.labels_to_names[label]))
self.objectLabelList.append(str(config_labels[label]))
self.objectListBox.insert(END, '(%d, %d) -> (%d, %d)' % (b[0], b[1], b[2], b[3]) + ': ' +
str(config.labels_to_names[label]))
str(config_labels[label])+' '+str(int(score*100))+'%'
+' '+ m_name)
self.objectListBox.itemconfig(len(self.bboxIdList) - 1,
fg=config.COLORS[(len(self.bboxIdList) - 1) % len(config.COLORS)])
self.processingLabel.config(text="Done ")
Expand Down