image_selector_for_images.py

'''
Author: Lee Wannacott
Github: https://github.com/LeeWannacott/
'''
from cv2 import cv2
import numpy as np
import math
import easygui
import sys
import natsort
import glob
import random


''' 
This program has been designed for use in a machine learning - data science context; to speed up image selection.

Program designed for reading in images from a folder of images resizing the original images down into a grid and allowing 
the user to select a span of images and tag the images with text. A text file containing the frame numbers and tagging 
is produced on each press of the space bar; New images are also fed into the grid on Space bar.

Esc = Exit the program.
Left mouse click = Select span of images.
Right mouse click = Undo image span selection.
Space bar = Go to next set of images and store tagged images + frame number in text file.
'''


class CreateGrid: pass
grid = CreateGrid()


# read in images.
files = glob.glob("*.jpg")

# sort files into numerical order.
files = natsort.natsorted(files, reverse=False)

# ************************************************************************
# variables that can be change by the user:

# Number of rows is changeable value. Default value is 7.
grid.number_of_rows = 7

# Variables containing default window size for the grid - default values assume 1080p.
grid.window_width = 1800
grid.window_height = 1050

# Original images from video that will be resized down to fit into grid.
grid.initial_img_width = 720
grid.initial_img_height = 522

# ************************************************************************

grid.cell_height= 0
grid.cell_width = 0
grid.number_of_cells= 0
grid.number_of_columns = 0
grid.image_resize_x = 0
grid.image_resize_y = 0

# Calulating how many images useful for last grid of images. 
frames_in_image_list = int(len(files))
print('Amount of images: ' + str(frames_in_image_list))

def recalculate_window_stuff():
    (_, _, temp_window_width, temp_window_height) = cv2.getWindowImageRect('image_selector_from_video')
    if not (temp_window_width == -1 or temp_window_height == -1):
        grid.window_width = temp_window_width
        grid.window_height = temp_window_height
    print(f"Window size: {grid.window_width}x{grid.window_height}")

    grid.cell_height = grid.window_height // grid.number_of_rows
    grid.cell_width = int(grid.cell_height * (grid.initial_img_width / grid.initial_img_height))

    # cell_aspect_ratio and image_aspect_ratio used for assert statement; for debugging.
    cell_aspect_ratio = grid.cell_width / grid.cell_height
    image_aspect_ratio = grid.initial_img_width / grid.initial_img_height
    assert abs(cell_aspect_ratio - image_aspect_ratio) < 0.01, f"cell_aspect_ratio={cell_aspect_ratio} \
        image_aspect_ratio={image_aspect_ratio}"

    grid.number_of_columns = grid.window_width // grid.cell_width
    grid.image_resize_x = grid.cell_width / grid.initial_img_width
    grid.image_resize_y = grid.cell_height / grid.initial_img_height
    grid.number_of_cells = grid.number_of_rows * grid.number_of_columns
    print(f"Cell width: {str(grid.cell_width)}, Cell Height: {str(grid.cell_height)}, "
            f"Number of rows: {str(grid.number_of_rows)},Number of Columns: {str(grid.number_of_columns)}")

class ClickFunctionality: pass
mouse_click = ClickFunctionality()
mouse_click.last_mouse_button_clicked = []
mouse_click.enable_draw_on_grid = False
mouse_click.coordinates = []


class ImageSelection: pass
image_selection = ImageSelection()
image_selection.drawn_one_cell_or_span = []
image_selection.tagged_as = ''
image_selection.new_image = ''
image_selection.image_list_temporary = []
image_selection.image_list = []
image_selection.cell_numbers_selection_temporary = []
image_selection.cell_numbers_selection_for_drawing_text = []

class CreateBoundingBox: pass
bounding_box = CreateBoundingBox()
bounding_box.bounding_box_start_coordinates_x_y = ()
bounding_box.list_of_bounding_box_coordinates = []
bounding_box.temp_list_cells_with_bboxes = []
bounding_box.temp_dict_and_cell_number_bboxes = {}
bounding_box.perm_dict_of_cell_num_and_bbox = {}

class CreateTextFile: pass
create_text_file = CreateTextFile()
create_text_file.frame_numbers_list = []
create_text_file.list_of_frames_to_keep = []
create_text_file.image_list_to_print = []
create_text_file.image_list_to_keep = []
create_text_file.cell_numbers_list_for_each_grid = []
def click_event(event, x, y, flags, param):
    # Mouse click for left button
    global mouseX
    global mouseY
    global global_click_event

    global_click_event = event
    mouseX = x
    mouseY = y
    if event == cv2.EVENT_LBUTTONDOWN and image_selection.tagged_as == '' \
        and mouse_click.enable_draw_on_grid == True:
        mouse_click.last_mouse_button_clicked.append('left')
        # Gets row and column number on left mouse click
        col_number = x / grid.cell_width
        x1 = math.trunc(col_number)
        row_number = y / grid.cell_height
        y1 = math.trunc(row_number)

        # Get cell number based on coordinates x1, y1
        def coordinate_to_cell(x, y):
            # https://stackoverflow.com/questions/9816024/coordinates-to-grid-box-number
            cell_number = int(x + (y * grid.number_of_columns))

            # Temporary list for drawing rectangles
            image_selection.cell_numbers_selection_temporary.append(cell_number)
            image_selection.cell_numbers_selection_for_drawing_text.append(cell_number)
            create_text_file.cell_numbers_list_for_each_grid.append(cell_number)
            return cell_number

        # For use in drawing rectangle function.
        cell = coordinate_to_cell(x1, y1)

        # Making copy of image for undo function.
        image_selection.new_image = param[0].copy()
        image_selection.image_list.append(image_selection.new_image)

        # Uses the cell number converts to coordinates and draws rectangles
        def draw_rectangles(cell_number):
            # !SO 8669189/converting-numbers-within-grid-to-their-corresponding-x-y-coordinates
            x2 = cell_number % grid.number_of_columns
            y2 = cell_number // grid.number_of_columns
            cell_x_position = x2 * grid.cell_width
            cell_y_position = y2 * grid.cell_height
            draw_to_x = cell_x_position + grid.cell_width
            draw_to_y = cell_y_position + grid.cell_height

            cv2.rectangle(param[0], (int(cell_x_position), int(cell_y_position)), (int(draw_to_x), int(draw_to_y)),
                          (152,251,152), 2)
            cv2.imshow('image_selector_from_video', param[0])

        # draw rectangles between two grid images
        def draw_rectangles_span():
            # Draws rectangle in the first cell
            if len(image_selection.cell_numbers_selection_temporary) == 1:
                draw_rectangles(cell)
                image_selection.drawn_one_cell_or_span.append('one')

            if len(image_selection.cell_numbers_selection_temporary) >= 2:
                between_backwards = list(
                    range(int(image_selection.cell_numbers_selection_temporary[-1]), \
                        int(image_selection.cell_numbers_selection_temporary[-2] + 1)))  # backwards
                between_forwards = list(
                    range(int(image_selection.cell_numbers_selection_temporary[-2]), \
                        int(image_selection.cell_numbers_selection_temporary[-1] + 1)))  # forwards

                # Draw rectangles backwards
                if image_selection.cell_numbers_selection_temporary[-1] < \
                     image_selection.cell_numbers_selection_temporary[-2]:
                    for next_cell1 in between_backwards:
                        draw_rectangles(next_cell1)

                # Draw rectangles forwards
                if image_selection.cell_numbers_selection_temporary[-1] > \
                     image_selection.cell_numbers_selection_temporary[-2]:
                    for next_cell2 in between_forwards:
                        draw_rectangles(next_cell2)

                # Clear temporary list so another two squares can be selected
                image_selection.cell_numbers_selection_temporary.clear()

                image_selection.drawn_one_cell_or_span.append('span')

                if image_selection.tagged_as == '':
                    # Window_open = True
                    image_selection.tagged_as = 'window open'
                    image_selection.tagged_as = easygui.enterbox("What is tagged in the images?")
                    # replace
                    if image_selection.tagged_as is not None:
                        image_selection.tagged_as = image_selection.tagged_as.replace(" ","")
                        draw_label_on_selection_span_of_images()
                        ####

                if image_selection.tagged_as != 'window open':
                    if image_selection.tagged_as is None:
                        image_selection.tagged_as = ''
                    image_selection.image_list_temporary.append(image_selection.tagged_as)

                    image_selection.tagged_as = ''

        def draw_label(cell_number):
            # !SO 8669189/converting-numbers-within-grid-to-their-corresponding-x-y-coordinates
            x2 = cell_number % grid.number_of_columns
            y2 = cell_number // grid.number_of_columns
            cell_x_position = x2 * grid.cell_width
            cell_y_position = y2 * grid.cell_height
            cv2.putText(img=param[0], text=image_selection.tagged_as, org=(cell_x_position+5, cell_y_position+20),
                        fontFace=cv2.FONT_ITALIC, fontScale=0.5, color=(255,105,180),
                        thickness=1, lineType=cv2.LINE_AA)
            cv2.imshow('image_selector_from_video', param[0])

        def draw_label_on_selection_span_of_images():
            if len(image_selection.cell_numbers_selection_for_drawing_text) >= 2:
                between_backwards = list(
                    range(int(image_selection.cell_numbers_selection_for_drawing_text[-1]),
                          int(image_selection.cell_numbers_selection_for_drawing_text[-2] + 1)))  # backwards
                between_forwards = list(
                    range(int(image_selection.cell_numbers_selection_for_drawing_text[-2]),
                          int(image_selection.cell_numbers_selection_for_drawing_text[-1] + 1)))  # forwards

                # Draw rectangles backwards
                if image_selection.cell_numbers_selection_for_drawing_text[-1] < \
                     image_selection.cell_numbers_selection_for_drawing_text[
                    -2]:
                    for backward_cell in between_backwards:
                        draw_label(backward_cell)

                # Draw rectangles forwards
                if image_selection.cell_numbers_selection_for_drawing_text[-1] > \
                     image_selection.cell_numbers_selection_for_drawing_text[
                    -2]:
                    for forward_cell in between_forwards:
                        draw_label(forward_cell)

        draw_rectangles_span()

    elif event == cv2.EVENT_MBUTTONUP and image_selection.tagged_as == '' \
        and mouse_click.enable_draw_on_grid == True and len(
            mouse_click.last_mouse_button_clicked) > 0 and \
                 mouse_click.last_mouse_button_clicked[-1] == 'left':


        # Allows undo if one image selected.
        if len(image_selection.image_list) >= 1 and image_selection.drawn_one_cell_or_span[-1] == 'one':
            param[0] = image_selection.image_list[-1]
            cv2.imshow('image_selector_from_video', image_selection.image_list[-1])
            cv2.waitKey(1)
            image_selection.image_list.pop()
            create_text_file.cell_numbers_list_for_each_grid.pop()
            mouse_click.last_mouse_button_clicked.pop()
            image_selection.drawn_one_cell_or_span.pop()
            image_selection.cell_numbers_selection_temporary.clear()

        # Allows undo if span of images selected.
        elif len(image_selection.image_list) >= 2 and image_selection.drawn_one_cell_or_span[-1] == 'span':
            param[0] = image_selection.image_list[-2]
            cv2.imshow('image_selector_from_video', image_selection.image_list[-2])
            cv2.waitKey(1)

            # Popping lists to undo selection span of images
            image_selection.image_list.pop()
            image_selection.image_list.pop()
            mouse_click.last_mouse_button_clicked.pop()
            mouse_click.last_mouse_button_clicked.pop()
            image_selection.drawn_one_cell_or_span.pop()
            image_selection.drawn_one_cell_or_span.pop()
            create_text_file.cell_numbers_list_for_each_grid.pop()
            create_text_file.cell_numbers_list_for_each_grid.pop()

            if len(image_selection.image_list_temporary) > 0:
                image_selection.image_list_temporary.pop()

    elif event == cv2.EVENT_MBUTTONUP and mouse_click.enable_draw_on_grid == True and \
        len(mouse_click.last_mouse_button_clicked) > 0 and \
             mouse_click.last_mouse_button_clicked[-1] == 'right':
        if len(image_selection.image_list) >= 1 and mouse_click.last_mouse_button_clicked[-1] == 'right':
            param[0] = image_selection.image_list[-1]
            cv2.imshow('image_selector_from_video', image_selection.image_list[-1])
            cv2.waitKey(1)

            image_selection.image_list.pop()
            mouse_click.last_mouse_button_clicked.pop()
            # Allows for the potential to have multiple boundary boxes in same cell
            if len(bounding_box.temp_dict_and_cell_number_bboxes[bounding_box. \
                temp_list_cells_with_bboxes[-1]]) > 4:

                # Removing four coordinates if more than one boundary box in cell.
                for i in range(4):
                    del bounding_box.temp_dict_and_cell_number_bboxes[bounding_box. \
                        temp_list_cells_with_bboxes[-1]][-1]

            # Removes boundary boxes if in a single cell
            elif len(bounding_box.temp_dict_and_cell_number_bboxes[bounding_box. \
                temp_list_cells_with_bboxes[-1]]) == 4:
                bounding_box.temp_dict_and_cell_number_bboxes.pop(bounding_box. \
                    temp_list_cells_with_bboxes[-1])


            # Removes cell number from temporary list
            bounding_box.temp_list_cells_with_bboxes.pop()

    # Allow bounding boxes to be places over images
    elif global_click_event == cv2.EVENT_RBUTTONDOWN and len(image_selection.drawn_one_cell_or_span) > 0 and \
         image_selection.drawn_one_cell_or_span[-1] == 'span':

        def get_bounding_box_start_coordinates(x, y):
            x_start_boundary = x
            y_start_boundary = y
            # Calculates starting cell number
            col_number = x / grid.cell_width
            x1 = math.trunc(col_number)
            row_number = y / grid.cell_height
            y1 = math.trunc(row_number)
            cell_number_on_start_of_drawing = int(x1 + (y1 * grid.number_of_columns))
            return x_start_boundary, y_start_boundary, cell_number_on_start_of_drawing

        bounding_box.bounding_box_start_coordinates_x_y = get_bounding_box_start_coordinates(x, y)

        # Handle visual drawing of placing bounding box; to give user feedback of where they are placing bbox.
        global allow_draw_bbox
        allow_draw_bbox = True
        while allow_draw_bbox == True:
            draw_bbox_images = param[0].copy()
            cv2.rectangle(draw_bbox_images, (bounding_box.bounding_box_start_coordinates_x_y[0], \
                 bounding_box.bounding_box_start_coordinates_x_y[1]),(mouseX, mouseY), (32,178,170), 2)
            cv2.imshow('image_selector_from_video',draw_bbox_images)
            cv2.waitKey(10)


    elif event == cv2.EVENT_RBUTTONUP and len(image_selection.drawn_one_cell_or_span) > 0 and \
         image_selection.drawn_one_cell_or_span[-1] == 'span':
        # Get cell number at end of bounding box
        col_number = x / grid.cell_width
        cell_x = math.trunc(col_number)
        row_number = y / grid.cell_height
        cell_y = math.trunc(row_number)
        cell_number_on_end_of_drawing = int(cell_x + (cell_y * grid.number_of_columns))
        # Minus x,y positions to get cells relative position
        cell_x_position = cell_x * grid.cell_width
        cell_y_position = cell_y * grid.cell_height
        allow_draw_bbox = False

        def draw_boundary_box(x, y, start_boundary_x_and_y):
            # Making copy for un-drawing bounding box
            new_image_boundary = param[0].copy()
            image_selection.image_list.append(new_image_boundary)

            # Draws bounding box
            end_boundary_x_and_y = (x, y)
            cv2.rectangle(param[0], (start_boundary_x_and_y[0], start_boundary_x_and_y[1]),
                          (end_boundary_x_and_y[0], end_boundary_x_and_y[1]), (0,255,127), 2)
            cv2.imshow('image_selector_from_video', param[0])

            # Used for un-drawing bounding box logic
            mouse_click.last_mouse_button_clicked.append('right')

            # Makes the coordinates relative to within that cell rather than the whole window.
            cell_start_relative_position_x = start_boundary_x_and_y[0] - cell_x_position
            cell_start_relative_position_y = start_boundary_x_and_y[1] - cell_y_position
            cell_end_relative_position_x = end_boundary_x_and_y[0] - cell_x_position
            cell_end_relative_position_y = end_boundary_x_and_y[1] - cell_y_position

            # Resizing coordinates back to the those of the original sized images rounded to whole pixels.
            cell_start_relative_position_x_resized = round(cell_start_relative_position_x / grid.image_resize_x)
            cell_start_relative_position_y_resized = round(cell_start_relative_position_y / grid.image_resize_y)
            cell_end_relative_position_x_resized = round(cell_end_relative_position_x / grid.image_resize_x)
            cell_end_relative_position_y_resized = round(cell_end_relative_position_y / grid.image_resize_y)

            # Gets bounding box x,y start and x,y end, relative to that individual cell.
            list_bounding_box_coordinates = [cell_start_relative_position_x_resized, cell_start_relative_position_y_resized, \
                cell_end_relative_position_x_resized, cell_end_relative_position_y_resized]

            # Cells the user have selected that contain bounding boxes
            bounding_box.temp_list_cells_with_bboxes.append(bounding_box.bounding_box_start_coordinates_x_y[2])




            # Checks if there is already a key from there already being a bounding box in the cell
            if bounding_box.bounding_box_start_coordinates_x_y[2] in bounding_box.temp_dict_and_cell_number_bboxes:
                bounding_box.temp_dict_and_cell_number_bboxes[bounding_box.bounding_box_start_coordinates_x_y[2]]. \
                    extend(list_bounding_box_coordinates)

            else:
                bounding_box.temp_dict_and_cell_number_bboxes[bounding_box.bounding_box_start_coordinates_x_y[2]] = \
                     list_bounding_box_coordinates

        # Checks if its still within the same cell and if so draws bounding box
        if bounding_box.bounding_box_start_coordinates_x_y[2] == cell_number_on_end_of_drawing and \ 
        len(create_text_file.cell_numbers_list_for_each_grid) > 0:
            # Check if drawing boundary boxes in span of images that has been selected.
            if bounding_box.bounding_box_start_coordinates_x_y[2] in range(create_text_file.cell_numbers_list_for_each_grid[-2], \
                 create_text_file.cell_numbers_list_for_each_grid[-1]+1)  \
            or bounding_box.bounding_box_start_coordinates_x_y[2] in range(create_text_file.cell_numbers_list_for_each_grid[-1], \
                 create_text_file.cell_numbers_list_for_each_grid[-2]+1):
                draw_boundary_box(x, y, bounding_box.bounding_box_start_coordinates_x_y)


index = 0

cv2.namedWindow('image_selector_from_video', cv2.WINDOW_NORMAL)
cv2.resizeWindow('image_selector_from_video', grid.window_width, grid.window_height)


def image_grid(index, x_offset=0, y_offset=0, i=0):
    recalculate_window_stuff()
    # Resetting large image/grid to black each time.
    l_img = np.zeros((grid.window_height, grid.window_width, 3), np.uint8)
    cv2.imshow('image_selector_from_video', l_img)
    cv2.waitKey(1)
    index_for_frame_list = index

    while i < grid.number_of_cells:

        mouse_click.enable_draw_on_grid = False
        for i, s_image in enumerate(files[index:]):

            s_image = cv2.imread(s_image)
            #  Resize image
            s_image = cv2.resize(s_image, (0, 0), None, grid.image_resize_x, grid.image_resize_y)

            # Put small images onto large image
            x_offset = (i % grid.number_of_columns) * int(grid.cell_width)
            y_offset = (i // grid.number_of_columns) * int(grid.cell_height)
            l_img[y_offset:y_offset + s_image.shape[0], x_offset:x_offset + s_image.shape[1]] = s_image

            # Show each small images drawn
            cv2.imshow('image_selector_from_video', l_img)
            cv2.waitKey(1)

            # i for count of images, index keeps track of where you are in frames
            i += 1
            index += 1
            if i == grid.number_of_cells or index == frames_in_image_list:
                mouse_click.enable_draw_on_grid = True
                # Parameters passed to mouse click function
                param = [l_img, index_for_frame_list]
                while True:

                    cv2.setMouseCallback('image_selector_from_video', click_event, param)
                    c = cv2.waitKey(1)

                    if c == 27:  # Escape
                        print('Esc pressed to Exit')
                        sys.exit()

                    elif c == 32:  # Space
                        print('Spacebar pressed')

                        # Getting rid of an uneven number of cells in lists
                        if len(image_selection.cell_numbers_selection_temporary) % 2 == 0:
                            pass
                        else:
                            image_selection.cell_numbers_selection_temporary.pop()
                        if len(create_text_file.cell_numbers_list_for_each_grid) % 2 == 0:
                            pass
                        else:
                            create_text_file.cell_numbers_list_for_each_grid.pop()

                        # Clears lists for the case of single selected rectangle
                        image_selection.cell_numbers_selection_temporary.clear()
                        image_selection.image_list.clear()

                        # Creates a image list from the temporary image lists of each grid of images.
                        def create_permanent_image_list():
                            for images in image_selection.image_list_temporary:
                                create_text_file.image_list_to_keep.append(images)

                        create_permanent_image_list()

                        # Transforms cell number in temporary grid into frame and appends to frame numbers list
                        def each_grid_cells_into_frames_list():
                            for cell in create_text_file.cell_numbers_list_for_each_grid:
                                frame_number = int(param[1]) + int(cell)
                                create_text_file.frame_numbers_list.append(frame_number)

                        each_grid_cells_into_frames_list()

                        # Calculates frames spans backwards and forwards
                        def make_list_of_frames_to_keep(image_count=0):
                            # Putting into a list of two's for calculating frame spans
                            frame_numbers_list_sliced = zip(create_text_file.frame_numbers_list[0::2],
                                                            create_text_file.frame_numbers_list[1::2])
                            # Clear lists when space pressed
                            create_text_file.list_of_frames_to_keep.clear()
                            create_text_file.image_list_to_print.clear()

                            for numbers in frame_numbers_list_sliced:
                                # Forward frame spans
                                if numbers[0] < numbers[1]:
                                    for frames1 in range(numbers[0], numbers[1] + 1):
                                        create_text_file.list_of_frames_to_keep.append(frames1)
                                        create_text_file.image_list_to_print.append(
                                            create_text_file.image_list_to_keep[image_count])
                                    image_count += 1

                                # Backward frame spans
                                elif numbers[0] > numbers[1]:
                                    for frames2 in range(numbers[1], numbers[0] + 1):
                                        create_text_file.list_of_frames_to_keep.append(frames2)
                                        create_text_file.image_list_to_print.append(
                                            create_text_file.image_list_to_keep[image_count])
                                    image_count += 1

                        make_list_of_frames_to_keep()

                        def bounding_box_cell_keys_to_frames():
                            # Changing the key of dict: cell numbers into frame numbers.
                            temp_dict_to_append = {k + int(param[1]): v for (k, v) in
                                                   bounding_box.temp_dict_and_cell_number_bboxes.items()}
                            bounding_box.perm_dict_of_cell_num_and_bbox.update(temp_dict_to_append)

                        bounding_box_cell_keys_to_frames()

                        # Clear temp dict of bounding boxes ready for next lot of frames to have bounding boxes.
                        bounding_box.temp_dict_and_cell_number_bboxes.clear()

                        # Clearing temporary lists, so ready for the next lot of images.
                        create_text_file.cell_numbers_list_for_each_grid.clear()
                        image_selection.image_list_temporary.clear()

                        def creates_text_file():
                            # Creating text file with frame numbers and what user tags images as.
                            file = open('List_of_images.txt', 'w')

                            for i, frame in enumerate(create_text_file.list_of_frames_to_keep):
                                if frame in bounding_box.perm_dict_of_cell_num_and_bbox.keys():
                                    file.write(
                                        f'{frame} {create_text_file.image_list_to_print[i]} \
                                            {bounding_box.perm_dict_of_cell_num_and_bbox[frame]} \n')
                                else:
                                    file.write(f'{frame} {create_text_file.image_list_to_print[i]} \n')

                            file.close()

                        # Output text file on each press of Space.
                        creates_text_file()

                        # Continue calling function if hasn't hit end of video
                        if frames_in_image_list != index:
                            return image_grid(index)

                        # If Space bar pressed and end of the video exit out of loop and save tagging.
                        elif (frames_in_image_list == index and index >
                              grid.number_of_cells):
                            creates_text_file()
                            return


# Calls function with index of 0.
image_grid(index)

# Closes all the windows.
cv2.destroyAllWindows()

# Shuts program down.
sys.exit()