Script to filter images from input directory to output directory

data_filter.py

# Run with: python data_filter.py --input_path <input path> --output_path <output path>
# Controls: z to copy image into output directory, x to skip to next image.

import pygame
import argparse
import shutil
import os
import argparse

def load_image(path):
    return pygame.image.load(path).convert()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_path', type=str, required=True)
    parser.add_argument('--output_path', type=str, required=True)
    args = parser.parse_args()

    screen_width = 500
    screen_height = 500
    pygame.init()
    screen = pygame.display.set_mode([screen_width, screen_height])
    pygame.display.set_caption('data filter')

    input_path = args.input_path
    output_path = args.output_path
    input_list = os.listdir(input_path)
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    
    def get_current_input_path():
        return os.path.join(input_path, input_list[current_image_id])

    def get_current_output_path():
        return os.path.join(output_path, input_list[current_image_id])

    current_image_id = 0
    current_input_path = get_current_input_path()
    current_image = load_image(current_input_path)
    
    next_input = False
    running = True
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                running = False
                break
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_z:
                    shutil.copy(current_input_path, get_current_output_path())
                    next_input = True
                if event.key == pygame.K_x:
                    next_input = True
                if next_input == True:
                    next_input = False
                    current_image_id += 1
                    current_input_path = get_current_input_path()
                    current_image = load_image(current_input_path)
        if running == False:
            break

        screen.fill([0, 0, 0])
        screen.blit(current_image, (0,0))
        pygame.display.flip()
        pygame.time.delay(20)

if __name__ == '__main__':
    main()