Simple image segmentation app

Question

I often segment images by hand to generate training data for machine learning models. To this end, I created a little TkInter app that displays images and allows me to export labeled polygons to JSON. I'm looking for feedback on the design of the project and coding style.

import tkinter as tk
from collections import defaultdict, namedtuple
from json import dump, load
from os import remove
from os.path import splitext, basename, join, isfile
from glob import glob
from PIL import ImageTk, Image
from argparse import ArgumentParser

point = namedtuple('point', 'x y')

class SimpleSegment(object):

    def __init__(self, image_path, label_directory, image_extension="png", lambda_func=None):
        self.files = sorted(glob(join(image_path, "*" + image_extension)))
        self.label_directory = label_directory
        self.drawing = False
        self.object_ids = defaultdict(list)
        self.poly_coords = defaultdict(list)
        self.poly_count = 0
        self.image_index = 0
        self.save_on_next = True
        root = tk.Tk()
        self.canvas = tk.Canvas(root)
        self.canvas.old_coords = None
        self.canvas.pack(fill='both', expand='yes')
        self._draw_image(self.files[self.image_index])
        self.info_box = tk.StringVar()
        self._update_string()
        self.lr_label = tk.Label(root, textvariable=self.info_box)
        self.lr_label.place(relx=1.0, rely=1.0, anchor='se')
        root.bind('<Motion>', self._draw_polygon)
        root.bind('<Button-3>', self._not)
        root.bind('u', self._remove)
        root.bind('j', self._prev_image)
        root.bind('l', self._skip_100)
        root.bind('h', self._back_100)
        root.bind('k', self._next_image)
        root.bind('d', self._delete_json_file)
        root.mainloop()

    def _skip_100(self, event):
        self.image_index = min(len(self.files)-1, self.image_index + 100)
        self._draw_image(self.files[self.image_index])
        self._update_string()

    def _back_100(self, event):
        self.image_index = max(0, self.image_index - 100)
        self._draw_image(self.files[self.image_index])
        self._update_string()

    def _update_string(self):
        string_to_format = 'right click to start and stop drawing. k: next image, j: prev image, h: skip 100 back. l: skip 100 forward. u: remove polygon. d: delete all labels. images: {} of {}'
        string = string_to_format.format(self.image_index+1, len(self.files))
        self.info_box.set(string)


    def _delete_json_file(self, event):
        json = self._create_json_filename(self.files[self.image_index])
        try:
            remove(json)
        except FileNotFoundError as e:
            print("Labels don't exist for this image")
        self.save_on_next = False

    def _draw_image(self, image_path):
        img = ImageTk.PhotoImage(Image.open(image_path))
        self.canvas.image = img
        self.canvas.create_image(0, 0, image=img, anchor='nw')

    def _create_json_filename(self, filename):
        out_filename = join(self.label_directory,
                splitext(basename(filename))[0]) + ".json"
        return out_filename

    def _dump_coord_dict(self):
        if not self.save_on_next:
            return
        out_filename = self._create_json_filename(self.files[self.image_index])
        with open(out_filename, 'w') as f:
            dump(self.poly_coords, f)
        self.object_ids = defaultdict(list)
        self.poly_coords = defaultdict(list)
        self.poly_count = 0

    def _next_image(self, event):
        if len(self.poly_coords):
            self._dump_coord_dict()
        self.image_index += 1
        if self.image_index > len(self.files) - 1:
            print("End of images!")
            return
        self._draw_image(self.files[self.image_index])
        self._load_and_draw_predrawn_polys(self.files[self.image_index])
        self._update_string()


    def _load_and_draw_predrawn_polys(self, image_filename):
        json_filename = self._create_json_filename(image_filename)
        if not isfile(json_filename):
            return
        with open(json_filename, 'r') as f:
            poly_coords = load(f)
        for poly_id in poly_coords:
            self.poly_count += 1 
            x1, y1 = poly_coords[poly_id][0], poly_coords[poly_id][1]
            for i, poly_coord in enumerate(poly_coords[poly_id][1:]):
                # reconstruct neccessary structures
                obj_id = self.canvas.create_line(poly_coord[0], poly_coord[1], x1, y1,
                        fill='red', width=2)
                x1, y1 = poly_coord[0], poly_coord[1]
                self.object_ids[self.poly_count].append(obj_id)
                self.poly_coords[self.poly_count].append(point(x=poly_coord[0], y=poly_coord[1]))

    def _prev_image(self, event):
        if self.image_index <= 0:
            print("Beginning of images!")
            return
        self.object_ids = defaultdict(list)
        self.poly_coords = defaultdict(list)
        self.poly_count = 0
        self.image_index -= 1
        self._draw_image(self.files[self.image_index])
        self._load_and_draw_predrawn_polys(self.files[self.image_index])
        self._update_string()

    def _not(self, event):
        if not self.drawing:
            self.canvas.old_coords = None
            self.drawing = True
            self.poly_count += 1
            self.save_on_next = True
        else:
            self.drawing = False

    def _remove(self, event):
        if self.poly_count == 0:
            return
        obj_ids = self.object_ids.pop(self.poly_count)
        self.poly_coords.pop(self.poly_count)
        for c in obj_ids:
            self.canvas.delete(c)
        self.poly_count -= 1

    def _draw_polygon(self, event):
        if self.drawing:
            x, y = event.x, event.y
            if self.canvas.old_coords:
                x1, y1 = self.canvas.old_coords
                obj_id = self.canvas.create_line(x, y, x1, y1, fill='red', width=2)
                self.object_ids[self.poly_count].append(obj_id)
                self.poly_coords[self.poly_count].append(point(x=x, y=y))
            self.canvas.old_coords = x, y

if __name__ == '__main__':

     ap = ArgumentParser()
     ap.add_argument('--image-directory', type=str, help='directory where the images are stored',
             required=True)
     ap.add_argument('--label-directory', type=str, help='directory to save JSON labels in',
             required=True)
     ap.add_argument('--image-extension', type=str, help='extension of images', default='png')
     args = ap.parse_args()
     SimpleSegment(args.image_directory, args.label_directory, image_extension=args.image_extension)

Answer 1

This is just how I like things, but I'd add a little more spacing in some parts. Your code is quite squished-together looking. For example, I'd write _delete_json_file closer to

def _delete_json_file(self, event):
    json = self._create_json_filename(self.files[self.image_index])

    try:
        remove(json)

    except FileNotFoundError as e:
        print("Labels don't exist for this image")

    self.save_on_next = False

and for _next_image

def _next_image(self, event):
    if len(self.poly_coords):
        self._dump_coord_dict()

    self.image_index += 1

    if self.image_index > len(self.files) - 1:
        print("End of images!")
        return

    self._draw_image(self.files[self.image_index])
    self._load_and_draw_predrawn_polys(self.files[self.image_index])
    self._update_string()

and

def _dump_coord_dict(self):
    if not self.save_on_next:
        return

    out_filename = self._create_json_filename(self.files[self.image_index])

    with open(out_filename, 'w') as f:
        dump(self.poly_coords, f)

    self.object_ids = defaultdict(list)
    self.poly_coords = defaultdict(list)
    self.poly_count = 0

It lengthens your code, but I've always found that this style makes code easier to read.

PEP8 does recommend the use of blank lines, to some extent:

Use blank lines in functions, sparingly, to indicate logical sections.

In your code above, I would consider a call to open with a context manager, or a try...except to be "logical section[s]".

---

And in _next_image, I don't think that early return half-way through the function reads nicely. I'd probably change that bit to something like

    if self.image_index > len(self.files) - 1:
        print("End of images!")

    else:
        self._draw_image(self.files[self.image_index])
        self._load_and_draw_predrawn_polys(self.files[self.image_index])
        self._update_string()

or, maybe flipped:

if self.image_index <= len(self.files) - 1:
    self._draw_image(self.files[self.image_index])
    self._load_and_draw_predrawn_polys(self.files[self.image_index])
    self._update_string()

else:
    print("End of images!")

I like putting the more complicated, bulkier code first unless there's a good reason not to (like negating the condition hurting readability)

---

point = namedtuple('point', 'x y')

point should be in capitals since it's really a class. I'd also go for typing.NamedTuple here too unless you really like the brevity of namedtuple:

from typing import NamedTuple

class Point(NamedTuple):
    x: int  # Or float. I can't offhand tell the type
    y: int

This gives the minor bonus of indicating types for the reader and IDE.

---

_not isn't a very descriptive name. I'd probably call it _toggle_drawing or something similar.

---

class SimpleSegment(object):
   . . .

Unless you really want backwards compatibility with Python 2, don't specify object as the base class. It's redundant.

---

. . . ].append(point(x=poly_coord[0], y=poly_coord[1]))

Assuming poly_coord is an iterable with exactly two elements, that can be written as

. . . ].append(point(*poly_coord))

Just "unpack" poly_coord into the call to the point constructor.