json_labeler.py

# TensorBoxPy3 https://github.com/SMH17/TensorBoxPy3

#A tool for hand labelling images generating json files
#
#Pass in the directory where you store your images and a filename, then select the points on the images.
#Every time you hit next a json object is generated.
#The clear button removes all selected points on the current image, undo cancel your last action.
#When all files in the directory are processed, the json file is written out.
#
#Is enforced convention that rects are in top left, bottom right order.

import sys
import json
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import matplotlib as mpl
import matplotlib.patches as patches
mpl.rcParams['toolbar'] = 'None'
from matplotlib.widgets import Button
from os import listdir
from os.path import isfile, join

json_images = []

top_corners = []
bottom_corners = []
patchCache = [] #the rectangles that get drawn on the image, stored so they can be removed in an orderly fashion

print("# TensorBoxPy3: target labeling")


def removeAllPatches():
    for patch in patchCache:
        patch.remove()
    patchCache[:] = []


def skip(event):  #called when the skip button is hit
    global filename
    if len(onlyfiles) == 0:
        outfile.write(json.dumps(json_images, indent = 1))
        plt.close()
    else:
        filename = path + "/" + onlyfiles.pop()
        image = mpimg.imread(filename)
        imshow_obj.set_data(image)
        top_corners[:] = []
        bottom_corners[:] = []
        removeAllPatches()


def next(event):  #called when the next button is hit
    global filename
    global json_images

    rects = []

    one_decimal = "{0:0.1f}"
    for i in range(len(top_corners)):
        x1 = float(one_decimal.format(top_corners[i][0]))
        x2 = float(one_decimal.format(bottom_corners[i][0]))
        y1 = float(one_decimal.format(top_corners[i][1]))
        y2 = float(one_decimal.format(bottom_corners[i][1]))

        #enforce x1,y1 = top left, x2,y2 = bottom right

        tlx = min(x1,x2)
        tly = min(y1,y2)
        brx = max(x1,x2)
        bry = max(y1,y2)

        bbox = dict([("x1",tlx),("y1",tly),("x2",brx),("y2",bry)])
        rects.append(bbox)

    json_image = dict([("image_path",filename),("rects",rects)])

    json_images.append(json_image)

    progress_outfile.write(json.dumps(json_image, indent = 1))

    if len(onlyfiles) == 0:
        outfile.write(json.dumps(json_images, indent = 1))
        plt.close()
    else:
        filename = path + "/" + onlyfiles.pop()
        image = mpimg.imread(filename)
        imshow_obj.set_data(image)
        top_corners[:] = []
        bottom_corners[:] = []
        removeAllPatches()


def clear(event): #called when the clear button is hit
    top_corners[:] = []
    bottom_corners[:] = []
    removeAllPatches()


def onclick(event):  #called when anywhere inside the window is clicked
    if event.xdata > 1 and event.ydata > 1:
        if (len(top_corners) > len(bottom_corners)):
            bottom_corners.append([event.xdata,event.ydata])
            patchCache.append(patches.Rectangle((top_corners[-1][0], top_corners[-1][1])
                                           ,bottom_corners[-1][0] - top_corners[-1][0], bottom_corners[-1][1] - top_corners[-1][1],
                                           hatch='/',fill=False))
            ax.add_patch(patchCache[-1])
            plt.draw()
        else:
            top_corners.append([event.xdata,event.ydata])


def undo(event):  #called when the undo button is hit
    # Only act when a path was drawn
    if (len(top_corners) ==  len(bottom_corners)):
        bottom_corners.pop()
        top_corners.pop()
        to_remove = patchCache.pop()
        to_remove.remove()

ax = plt.gca()

#get our files for processing
if len(sys.argv) < 3:
    print("Too few params, try something like:  python json_labeler.py trainfolder640x480 train.json")
    exit()
path = sys.argv[1]
outfile_name = sys.argv[2]
outfile = open(outfile_name, 'w')
progress_outfile = open(outfile_name + "_work", 'w')

onlyfiles = [f for f in listdir(path) if isfile(join(path, f))]

filename = path + "/" + onlyfiles.pop()
image = mpimg.imread(filename)
imshow_obj = ax.imshow(image)

plt.axis("off")
fig = plt.gcf()
fig.canvas.mpl_connect('button_press_event', onclick)

#add the buttons to the bottom of the window
axundo = plt.axes([0.59, 0.01, 0.1, 0.075])
axnext = plt.axes([0.7, 0.01, 0.1, 0.075])
axclear = plt.axes([0.81, 0.01, 0.1, 0.075])
axskip = plt.axes([0.92, 0.01, 0.1, 0.075])
bundo = Button(axundo, 'Undo')
bundo.on_clicked(undo)
bnext = Button(axnext, 'Next')
bnext.on_clicked(next)
bclear = Button(axclear, 'Clear')
bclear.on_clicked(clear)
bskip = Button(axskip, 'Skip')
bskip.on_clicked(skip)
plt.show()

outfile.close()
progress_outfile.close()

print("labelling finished")