diff --git a/VisioNomicon/__init__.py b/VisioNomicon/__init__.py index 71777f1..cd15240 100644 --- a/VisioNomicon/__init__.py +++ b/VisioNomicon/__init__.py @@ -1,8 +1,8 @@ """ VisioNomicon. -A utility leveraging GPT-4V for image file renaming based on content. +A utility leveraging GPT-4V for image file renaming based on content. """ __version__ = "0.1.2" -__author__ = 'Rehan Rana' +__author__ = "Rehan Rana" diff --git a/VisioNomicon/args_handler.py b/VisioNomicon/args_handler.py index e11ef00..23c5f6b 100644 --- a/VisioNomicon/args_handler.py +++ b/VisioNomicon/args_handler.py @@ -1,29 +1,89 @@ -import argparse, os, sys +import argparse +import os NO_VAL = object() + def parse_cli_args(): - parser = argparse.ArgumentParser(description='A utility leveraging GPT-4V for image file renaming based on content.') - parser.add_argument('-f', '--files', type=str, nargs='*', help='Specify file paths of the images to create mapping for') - parser.add_argument('-o', '--output', type=str, nargs='?', help='Provide file path for the JSON mapping file') - parser.add_argument('-x', '--execute', type=str, nargs='?', help='Execute on given mapping', const=NO_VAL) - parser.add_argument('-ox', '--mapex', type=str, nargs='?', help='Map and execute on mapping at given location', const=NO_VAL) - parser.add_argument('-u', '--undo', type=str, nargs='?', help='Undo given mapping', const=NO_VAL) - parser.add_argument('-t', '--template', type=str, nargs='?', help='Define the template for renaming image files, without file extension', default='[SubjectDescription]_[MainColor/ColorScheme]_[StyleOrFeel]_[CompositionElement]') - parser.add_argument('-e', '--validation-retries', type=int, help='Number of times to retry if validation not passed', default=3) - parser.add_argument('-v', '--error-retries', type=int, help='Number of times to retry if error response recieved from OpenAI', default=3) - parser.add_argument('-E', '--ignore-validation-fail', action='store_true', help='If validation retries limit is reached, map file to original name instead of returning an error') - parser.add_argument('-V', '--ignore-error-fail', action='store_true', help='If error retries limit is reached, map file to original name instead of returning an error') + parser = argparse.ArgumentParser( + description="A utility leveraging GPT-4V for image file renaming based on content." + ) + parser.add_argument( + "-f", + "--files", + type=str, + nargs="*", + help="Specify file paths of the images to create mapping for", + ) + parser.add_argument( + "-o", + "--output", + type=str, + nargs="?", + help="Provide file path for the JSON mapping file", + ) + parser.add_argument( + "-x", + "--execute", + type=str, + nargs="?", + help="Execute on given mapping", + const=NO_VAL, + ) + parser.add_argument( + "-ox", + "--mapex", + type=str, + nargs="?", + help="Map and execute on mapping at given location", + const=NO_VAL, + ) + parser.add_argument( + "-u", "--undo", type=str, nargs="?", help="Undo given mapping", const=NO_VAL + ) + parser.add_argument( + "-t", + "--template", + type=str, + nargs="?", + help="Define the template for renaming image files, without file extension", + default="[SubjectDescription]_[MainColor/ColorScheme]_[StyleOrFeel]_[CompositionElement]", + ) + parser.add_argument( + "-e", + "--validation-retries", + type=int, + help="Number of times to retry if validation not passed", + default=3, + ) + parser.add_argument( + "-v", + "--error-retries", + type=int, + help="Number of times to retry if error response recieved from OpenAI", + default=3, + ) + parser.add_argument( + "-E", + "--ignore-validation-fail", + action="store_true", + help="If validation retries limit is reached, map file to original name instead of returning an error", + ) + parser.add_argument( + "-V", + "--ignore-error-fail", + action="store_true", + help="If error retries limit is reached, map file to original name instead of returning an error", + ) # if flag with value, equals value # if flag with no value, equals const value # if flag not used, equals None - # capture initial defaults before parsing # this is for the below 'hack' defaults = {action.dest: action.default for action in parser._actions} - + args = parser.parse_args() args_dict = vars(args) @@ -34,37 +94,39 @@ def parse_cli_args(): if args.undo is not None: # Check if any other arg changed from default non_default_args = [arg for arg in args_dict if args_dict[arg] != defaults[arg]] - + # Remove checked key since we don't need to check it against itself - non_default_args.remove('undo') - + non_default_args.remove("undo") + # If any other arguments changed, error if non_default_args: - parser.error('-u/--undo must not be used with any other arguments.') + parser.error("-u/--undo must not be used with any other arguments.") #################################################################################### if args.files is not None and len(args.files) == 0: parser.error("-f/--files requires a value") if args.output is not None and args.execute is not None: - parser.error("instead of using -o/--output along with -x/--execute, use -ox/--mapex") - + parser.error( + "instead of using -o/--output along with -x/--execute, use -ox/--mapex" + ) + if args.mapex is not None: if args.output is not None or args.execute is not None: - parser.error("-ox/--mapex should be used without -o/--output or -x/--execute") + parser.error( + "-ox/--mapex should be used without -o/--output or -x/--execute" + ) args.output = args.mapex args.execute = args.mapex - if args.output is not None and args.files is None: - parser.error('-o/--output must be used with -f/--files') + parser.error("-o/--output must be used with -f/--files") if args.template is None: - parser.error('used -t/--template with no value') - - supported_ext = ['.png', '.jpeg', '.jpg', '.webp', '.gif'] + parser.error("used -t/--template with no value") + supported_ext = [".png", ".jpeg", ".jpg", ".webp", ".gif"] # # get absolute paths where we need them @@ -83,7 +145,7 @@ def parse_cli_args(): for image_path in clean_paths: _, image_ext = os.path.splitext(image_path) if image_ext not in supported_ext: - parser.error('Filetype {} not supported'.format(image_ext)) + parser.error("Filetype {} not supported".format(image_ext)) args.files = clean_paths if args.output is not None and args.output != NO_VAL: @@ -94,5 +156,5 @@ def parse_cli_args(): if args.undo is not None and args.undo != NO_VAL: args.undo = os.path.abspath(args.undo) - + return args diff --git a/VisioNomicon/gpt.py b/VisioNomicon/gpt.py index 6e7e577..c29268c 100644 --- a/VisioNomicon/gpt.py +++ b/VisioNomicon/gpt.py @@ -1,13 +1,16 @@ from openai import OpenAI from pathlib import Path -import os, requests, base64, sys +import os +import requests +import base64 +import sys API_KEY = "" def set_api_key(): global API_KEY - not "OPENAI_API_KEY" in os.environ and sys.exit( + "OPENAI_API_KEY" not in os.environ and sys.exit( "Open AI API key not set. Set it using the OPENAI_API_KEY environment variable" ) API_KEY = os.environ.get("OPENAI_API_KEY") if API_KEY == "" else API_KEY diff --git a/VisioNomicon/main.py b/VisioNomicon/main.py index e326fce..6c63534 100644 --- a/VisioNomicon/main.py +++ b/VisioNomicon/main.py @@ -1,122 +1,152 @@ -import os, json, copy, glob -from .args_handler import * -from .gpt import * +import os +import json +import copy +import glob +from VisioNomicon.args_handler import * +from VisioNomicon.gpt import * from datetime import datetime DATA_DIR = "" + def main(): - # get data dir - global DATA_DIR - DATA_DIR = (os.environ.get("XDG_DATA_HOME") if "XDG_DATA_HOME" in os.environ else os.path.abspath("~/.local/share")) + "/visionomicon/" + # get data dir + global DATA_DIR + DATA_DIR = ( + os.environ.get("XDG_DATA_HOME") + if "XDG_DATA_HOME" in os.environ + else os.path.abspath("~/.local/share") + ) + "/visionomicon/" + + # make data dir if doesn't exist + not os.path.exists(DATA_DIR) and os.makedirs(DATA_DIR) - # make data dir if doesn't exist - not os.path.exists(DATA_DIR) and os.makedirs(DATA_DIR) + args = parse_cli_args() - args = parse_cli_args() + # if creating mapping + if args.files is not None: + new_filepaths: list[str] = generate_mapping(args) - # if creating mapping - if args.files is not None: - new_filepaths: list[str] = generate_mapping(args) + # have new and old, put them together into a json and save + save_mapping(args, new_filepaths) - # have new and old, put them together into a json and save - save_mapping(args, new_filepaths) + # if executing or undoing + if args.undo or args.execute: + rel_mapping_fp = args.execute if args.execute else args.undo + rename_from_mapping(rel_mapping_fp, args.undo is not None) - # if executing or undoing - if args.undo or args.execute: - rel_mapping_fp = args.execute if args.execute else args.undo - rename_from_mapping(rel_mapping_fp, args.undo is not None) def rename_from_mapping(rel_mapping_fp: str, undo: bool = False): - mapping_fp = get_mapping_name(rel_mapping_fp) - og_filepaths: list[str] = [] - new_filepaths: list[str] = [] - - with open(mapping_fp) as f: - data = json.load(f) - og_filepaths += list(data.keys()) - new_filepaths += list(data.values()) - - from_fps, to_fps = (new_filepaths, og_filepaths) if undo else (og_filepaths, new_filepaths) - - for i in range(len(from_fps)): - _, filename = os.path.split(to_fps[i]) - print("Renaming {} to {}".format(from_fps[i], '.../' + filename)) - os.rename(from_fps[i], to_fps[i]) + mapping_fp = get_mapping_name(rel_mapping_fp) + og_filepaths: list[str] = [] + new_filepaths: list[str] = [] + + with open(mapping_fp) as f: + data = json.load(f) + og_filepaths += list(data.keys()) + new_filepaths += list(data.values()) + + from_fps, to_fps = ( + (new_filepaths, og_filepaths) if undo else (og_filepaths, new_filepaths) + ) + + for i in range(len(from_fps)): + _, filename = os.path.split(to_fps[i]) + print("Renaming {} to {}".format(from_fps[i], ".../" + filename)) + os.rename(from_fps[i], to_fps[i]) + def get_mapping_name(cli_fp: str): - if cli_fp != NO_VAL: - return cli_fp - else: - # Join the directory with the file pattern - file_pattern = os.path.join(DATA_DIR, '*.json') + if cli_fp != NO_VAL: + return cli_fp + else: + # Join the directory with the file pattern + file_pattern = os.path.join(DATA_DIR, "*.json") - # Get list of files matching the file pattern - files = glob.glob(file_pattern) + # Get list of files matching the file pattern + files = glob.glob(file_pattern) + + # Sort files by creation time and get the last one + mapping = max(files, key=os.path.getctime) + return mapping - # Sort files by creation time and get the last one - mapping = max(files, key=os.path.getctime) - return mapping def save_mapping(args, new_filepaths: list[str]): - og_filepaths: list[str] = args.files - data = dict(zip(og_filepaths, new_filepaths)) - # data ready to dump, need to get mapping filename - mapping_filename = generate_mapping_name(args) + og_filepaths: list[str] = args.files + data = dict(zip(og_filepaths, new_filepaths)) + # data ready to dump, need to get mapping filename + mapping_filename = generate_mapping_name(args) + + with open(mapping_filename, "w") as file: + json.dump(data, file, indent=4) - with open(mapping_filename, 'w') as file: - json.dump(data, file, indent=4) def generate_mapping_name(args) -> str: - return args.output if args.output != NO_VAL else DATA_DIR + datetime.now().strftime("mapping-%Y-%m-%d-%H-%M-%S.json") - + return ( + args.output + if args.output != NO_VAL + else DATA_DIR + datetime.now().strftime("mapping-%Y-%m-%d-%H-%M-%S.json") + ) + + def generate_mapping(args) -> list[str]: - og_filepaths: list[str] = args.files - new_filepaths: list[str] = copy.deepcopy(og_filepaths) - - for i in range(len(new_filepaths)): - slicepoint = new_filepaths[i].rindex("/") + 1 - new_filepaths[i] = new_filepaths[i][:slicepoint] - - for i in range(len(og_filepaths)): - image_path = og_filepaths[i] - for j in range(args.validation_retries + 1): - print("Generating name...") - new_name = image_to_name(image_path, args) - print("Generated name {}".format(new_name)) - - _, image_ext = os.path.splitext(image_path) - new_filename = new_name + image_ext - new_fp = new_filepaths[i] + new_filename - - # if new_fp == image_path, that means image_to_name errored past retry limit - # mapping the file to the exact same name, keeping it the same - # this means it would not follow the template and fail validation, so we skip - if new_fp == image_path: - break - elif name_validation(new_name, args.template): - print("Name validated".format(new_name)) - break - elif j == args.validation_retries: - if args.ignore_validation_fail: - print("Failed validation {} time(s), leaving filename unchanged".format(args.validation_retries + 1)) - new_fp = image_path - else: - sys.exit("Failed validation {} time(s), aborting...".format(args.validation_retries + 1)) - else: - print("Generated name failed validation, regenerating...") - - new_filename_suffixed = new_filename - num_suffix = 1 - while new_fp in new_filepaths: - new_filename_suffixed = new_name + f"_{num_suffix}" + image_ext - new_fp = new_filepaths[i] + new_filename_suffixed - num_suffix += 1 - - new_filepaths[i] = new_fp - - print("File {} mapped to name {}\n".format(og_filepaths[i], new_filename_suffixed)) - return new_filepaths + og_filepaths: list[str] = args.files + new_filepaths: list[str] = copy.deepcopy(og_filepaths) + + for i in range(len(new_filepaths)): + slicepoint = new_filepaths[i].rindex("/") + 1 + new_filepaths[i] = new_filepaths[i][:slicepoint] + + for i in range(len(og_filepaths)): + image_path = og_filepaths[i] + for j in range(args.validation_retries + 1): + print("Generating name...") + new_name = image_to_name(image_path, args) + print("Generated name {}".format(new_name)) + + _, image_ext = os.path.splitext(image_path) + new_filename = new_name + image_ext + new_fp = new_filepaths[i] + new_filename + + # if new_fp == image_path, that means image_to_name errored past retry limit + # mapping the file to the exact same name, keeping it the same + # this means it would not follow the template and fail validation, so we skip + if new_fp == image_path: + break + elif name_validation(new_name, args.template): + print("Name validated".format()) + break + elif j == args.validation_retries: + if args.ignore_validation_fail: + print( + "Failed validation {} time(s), leaving filename unchanged".format( + args.validation_retries + 1 + ) + ) + new_fp = image_path + else: + sys.exit( + "Failed validation {} time(s), aborting...".format( + args.validation_retries + 1 + ) + ) + else: + print("Generated name failed validation, regenerating...") + + new_filename_suffixed = new_filename + num_suffix = 1 + while new_fp in new_filepaths: + new_filename_suffixed = new_name + f"_{num_suffix}" + image_ext + new_fp = new_filepaths[i] + new_filename_suffixed + num_suffix += 1 + + new_filepaths[i] = new_fp + + print( + "File {} mapped to name {}\n".format(og_filepaths[i], new_filename_suffixed) + ) + return new_filepaths + if __name__ == "__main__": main()