meme-classifier/main.py

107 lines
3.4 KiB
Python
Executable file

# Python project to classify memes
import os
import subprocess
import argparse
import time
import shutil
def main():
parser = argparse.ArgumentParser()
parser.add_argument("INPUT", help="folder with input files", type=str)
parser.add_argument("OUTPUT", help="folder for output", type=str)
parser.add_argument("-n", "--name", action="store_true")
parser.add_argument("-c", "--classify", action="store_true")
parser.add_argument(
"-r",
"--reformat",
action="store_true",
help="automatically convert videos to mp4 using ffmpeg",
)
parser.add_argument("-v", "--verbose", action="store_true")
args = parser.parse_args()
inputFolder: str = args.INPUT
if not inputFolder.endswith("/"):
inputFolder = inputFolder + "/"
outputFolder: str = args.OUTPUT
if not outputFolder.endswith("/"):
outputFolder = outputFolder + "/"
if not os.path.isdir(inputFolder):
print(f"\033[91mERROR: {inputFolder} does not exist")
exit(1)
if not os.path.isdir(outputFolder):
query = input(f"{outputFolder} does not exist. Do you want to create it? [Y/n]")
yes: tuple(str) = "yes", "y", ""
no: tuple(str) = "no", "n"
if query.lower() in yes:
if args.verbose:
print(f"Creating folder {outputFolder}")
os.mkdir(outputFolder)
elif query.lower() in no:
print("\033[91mNo output folder, aborting...")
exit(1)
files: [str] = os.listdir(inputFolder)
for index, file in enumerate(files):
if os.path.isdir(inputFolder + file):
files.pop(index)
if args.verbose:
for file in files:
print(file)
time.sleep(0.0001)
print(f"\nFound {len(files)} files to process")
for id, file in enumerate(files):
print(f"Processing file {id + 1}/{len(files)}")
subprocess.run(["mpv", f"{inputFolder}{file}"])
if args.name:
print(f"Old name: {file}")
extension: str = file.rpartition(".")[-1] or ""
if extension != "":
extension = "." + extension
newname = input("New name: ")
else:
newname = file
if args.classify:
print(f"Existing categories: {os.listdir(outputFolder)}")
category: str = input("Category: ")
if category == "delete":
os.remove(inputFolder + file)
elif not os.path.exists(outputFolder + category) or not os.path.isdir(
outputFolder + category
):
os.mkdir(outputFolder + category)
else:
category: str = ""
path: str = outputFolder + category + "/" + newname + extension
if os.path.exists(path):
newname = input("Name already in use, please choose another: ")
path = outputFolder + category + "/" + newname + extension
shutil.move(inputFolder + file, path)
if args.reformat:
if extension != ".mp4":
subprocess.run(
["ffmpeg"],
["-i"],
[path],
["-o"],
["mp4"],
[f"{outputFolder}{category}{newname}.mp4"],
)
if args.verbose:
print(f"converted {newname} from {extension} to .mp4")
if __name__ == "__main__":
main()