aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--.gitignore1
-rw-r--r--__init__.py0
-rw-r--r--file-tagger.py153
-rw-r--r--gui.py80
4 files changed, 234 insertions, 0 deletions
diff --git a/.gitignore b/.gitignore
index a20f059..c6cad8d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+/test/
# Created by https://www.toptal.com/developers/gitignore/api/visualstudiocode,linux,windows,macos,python,pycharm
# Edit at https://www.toptal.com/developers/gitignore?templates=visualstudiocode,linux,windows,macos,python,pycharm
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/__init__.py
diff --git a/file-tagger.py b/file-tagger.py
new file mode 100644
index 0000000..71a64e3
--- /dev/null
+++ b/file-tagger.py
@@ -0,0 +1,153 @@
+import numpy as np
+import argparse
+import os, sys
+from gui import GuiMain, GuiImage
+import cv2
+import logging
+import magic
+from subprocess import Popen, PIPE
+import re
+
+def dir_path(string):
+ if os.path.isdir(string):
+ return string
+ else:
+ raise NotADirectoryError(string)
+
+def tmsu_init(base):
+ logger = logging.getLogger(__name__)
+ if not os.path.exists(os.path.join(base, ".tmsu")):
+ logger.info("TMSU database does not exist, creating ...")
+ proc = Popen(["tmsu", "init"], cwd=base)
+ proc.wait()
+ logger.debug("TMSU returncode: {}".format(proc.returncode))
+ if proc.returncode != 0:
+ logger.error("Could not initialize TMSU database.")
+ return False
+ return True
+
+def tmsu_tags(base, file):
+ logger = logging.getLogger(__name__)
+ logger.debug("Getting existing tags for file {}".format(file))
+ tags = set()
+ proc = Popen(["tmsu", "tags", file], cwd=base, stdout=PIPE, stderr=PIPE)
+ proc.wait()
+ logger.debug("TMSU returncode: {}".format(proc.returncode))
+ if proc.returncode == 0:
+ tags.update(re.split("\s", proc.stdout.read().decode())[1:-1])
+ else:
+ logger.error("Could not get tags for file {}".format(file))
+ return tags
+
+def tmsu_tag(base, file, tags, untag=True):
+ logger = logging.getLogger(__name__)
+ if untag:
+ logger.debug("Untagging file")
+ proc = Popen(["tmsu", "untag", "--all", file], cwd=base, stdout=PIPE, stderr=PIPE)
+ proc.wait()
+ if proc.returncode != 0:
+ logger.error("Could not untag file {}".format(file))
+ if tags:
+ logger.debug("Writing tags {}".format(tags))
+ proc = Popen(["tmsu", "tag", file] + list(tags), cwd=base, stdout=PIPE, stderr=PIPE)
+ proc.wait()
+ if proc.returncode != 0:
+ logger.error("Could not write tags to file {}".format(file))
+ else:
+ logger.info("Tags are empty, ignoring")
+
+def walk(args):
+ logger = logging.getLogger(__name__)
+ logger.info("Walking files ...")
+ mime = magic.Magic(mime=True)
+ files = [os.path.abspath(os.path.join(dp, f)) for dp, dn, filenames in os.walk(args["base"]) for f in filenames]
+ logger.debug("Files: {}".format(files))
+ logger.info("Number of files found: {}".format(len(files)))
+
+ if args["predict_images"]:
+ from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input, decode_predictions
+ from tensorflow.keras.preprocessing import image
+ from tensorflow.keras.models import Model
+ model = ResNet50(weights="imagenet")
+
+ for file_path in files:
+ logger.info("Handling file {}".format(file_path))
+ tags = tmsu_tags(args["base"], file_path)
+ not_empty = bool(tags)
+ logger.info("Existing tags: {}".format(tags))
+ mime_type = mime.from_file(file_path)
+ if mime_type.split("/")[0] == "image":
+ logger.debug("File is image")
+ img = cv2.imread(file_path)
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+ img = cv2.resize(img, dsize=(800, 800), interpolation=cv2.INTER_CUBIC)
+ while(True):
+ if args["predict_images"]:
+ logger.info("Predicting image tags ...")
+ array = cv2.resize(img, dsize=(224, 224), interpolation=cv2.INTER_CUBIC)
+ array = np.expand_dims(array, axis=0)
+ array = preprocess_input(array)
+ predictions = model.predict(array)
+ classes = decode_predictions(predictions, top=10)
+ logger.debug("Predicted image classes: {}".format(classes[0]))
+ tags.update([name for _, name, _ in classes[0]])
+ logger.info("Predicted tags: {}".format(tags))
+ if args["gui_images"]:
+ logger.debug("Showing image GUI ...")
+ ret = GuiImage(img, tags).loop()
+ tags = set(ret[1]).difference({''})
+ if ret[0] == GuiImage.RETURN_ROTATE_90_CLOCKWISE:
+ img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
+ elif ret[0] == GuiImage.RETURN_ROTATE_90_COUNTERCLOCKWISE:
+ img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
+ elif ret[0] == GuiImage.RETURN_NEXT:
+ break
+ elif ret[0] == GuiImage.RETURN_ABORT:
+ return
+ continue
+ break
+ logger.info("Tagging {}".format(tags))
+ tmsu_tag(args["base"], file_path, tags, untag=not_empty)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description='Tag multiple files using TMSU.')
+ parser.add_argument('-b', '--base', nargs='?', default='./test', type=dir_path, help='Base directory for walking (default: %(default)s)')
+ parser.add_argument('-g', '--gui', nargs='?', const=1, default=False, type=bool, help='Show main GUI (default: %(default)s)')
+ parser.add_argument('--predict-images', nargs='?', const=1, default=False, type=bool, help='Use prediction for image tagging (default: %(default)s)')
+ parser.add_argument('--gui-images', nargs='?', const=1, default=False, type=bool, help='Show GUI for image tagging (default: %(default)s)')
+ parser.add_argument('--gui-audio', nargs='?', const=1, default=False, type=bool, help='Show GUI for audio tagging (default: %(default)s)')
+ parser.add_argument('--gui-video', nargs='?', const=1, default=False, type=bool, help='Show GUI for video tagging (default: %(default)s)')
+ parser.add_argument('--open-all', nargs='?', const=1, default=False, type=bool, help='Open all files with system default (default: %(default)s)')
+ parser.add_argument('-v', '--verbose', action="count", default=0, help="Verbosity level")
+ args = parser.parse_args()
+
+ if args.verbose == 0:
+ log_level = logging.WARNING
+ elif args.verbose == 1:
+ log_level = logging.INFO
+ elif args.verbose >= 2:
+ log_level = logging.DEBUG
+
+ logging.basicConfig(stream=sys.stdout, level=log_level)
+ logger = logging.getLogger(__name__)
+
+ args = {
+ "base": args.base,
+ "gui": args.gui,
+ "predict_images": args.predict_images,
+ "gui_images": args.gui_images,
+ "gui_audio": args.gui_audio,
+ "gui_video": args.gui_video,
+ "open_all": args.open_all,
+ "verbosity": args.verbose
+ }
+
+ logger.debug("args = {}".format(args))
+
+ if args["gui"]:
+ logger.debug("Starting main GUI ...")
+ args = GuiMain(args).loop()
+
+ if tmsu_init(args["base"]):
+ walk(args)
diff --git a/gui.py b/gui.py
new file mode 100644
index 0000000..cb30534
--- /dev/null
+++ b/gui.py
@@ -0,0 +1,80 @@
+from tkinter import *
+from tkinter import filedialog
+from PIL import Image, ImageTk
+import logging
+from enum import Enum
+
+class GuiMain(object):
+
+ def __init__(self, args):
+ self.__master = Tk()
+ self.__args = args
+ self.__base = StringVar(self.__master, value=args["base"])
+ self.__predict_images = BooleanVar(self.__master, value=args["predict_images"])
+ self.__gui_images = BooleanVar(self.__master, value=args["gui_images"])
+ self.__gui_audio = BooleanVar(self.__master, value=args["gui_audio"])
+ self.__gui_video = BooleanVar(self.__master, value=args["gui_video"])
+ self.__open_all = BooleanVar(self.__master, value=args["open_all"])
+
+ Label(self.__master, text="Base directory for walking:").grid(row=0, column=0)
+ Entry(self.__master, textvariable=self.__base).grid(row=0, column=1)
+ Button(self.__master, text="Browse", command=lambda: self.__browse(base)).grid(row=0, column=3)
+ Checkbutton(self.__master, text="Use prediction for image tagging", variable=self.__predict_images).grid(row=1, sticky=W)
+ Checkbutton(self.__master, text="Show GUI for image tagging", variable=self.__gui_images).grid(row=2, sticky=W)
+ Checkbutton(self.__master, text="Show GUI for audio tagging", variable=self.__gui_audio).grid(row=3, sticky=W)
+ Checkbutton(self.__master, text="Show GUI for video tagging", variable=self.__gui_video).grid(row=4, sticky=W)
+ Checkbutton(self.__master, text="Open all files with system default", variable=self.__open_all).grid(row=5, sticky=W)
+ Button(self.__master, text="Start", command=self.__master.destroy).grid(row=6)
+
+ def loop(self):
+ self.__master.mainloop()
+
+ self.__args["base"] = self.__base.get()
+ self.__args["predict_images"] = self.__predict_images.get()
+ self.__args["gui_images"] = self.__gui_images.get()
+ self.__args["gui_audio"] = self.__gui_audio.get()
+ self.__args["gui_video"] = self.__gui_video.get()
+ self.__args["open_all"] = self.__open_all.get()
+ return self.__args
+
+ def __browse(self, folder_path):
+ filename = filedialog.askdirectory()
+ folder_path.set(filename)
+
+class GuiImage(object):
+ RETURN_NEXT = 0,
+ RETURN_ROTATE_90_COUNTERCLOCKWISE = 1,
+ RETURN_ROTATE_90_CLOCKWISE = 2,
+ RETURN_ABORT = 3
+
+ def __init__(self, img, tags):
+ self.__ret = self.RETURN_NEXT
+ self.__master = Tk()
+ self.__tags = StringVar(self.__master, value=','.join(tags))
+ self.__image = ImageTk.PhotoImage(image=Image.fromarray(img).convert('RGB'))
+ Label(self.__master, width=800, height=800, image=self.__image).grid(row=0, column=0, columnspan=4)
+ Entry(self.__master, textvariable=self.__tags).grid(row=1, column=0, columnspan=4, sticky="we")
+ Button(self.__master, text="↺", command=self.__handle_rotate_90_counterclockwise).grid(row=2, column=0)
+ Button(self.__master, text="↻", command=self.__handle_rotate_90_clockwise).grid(row=2, column=1)
+ Button(self.__master, text="Next", command=self.__handle_next).grid(row=2, column=2)
+ Button(self.__master, text="Abort", command=self.__handle_abort).grid(row=2, column=3)
+
+ def loop(self):
+ self.__master.mainloop()
+ return (self.__ret, self.__tags.get().split(","))
+
+ def __handle_rotate_90_counterclockwise(self):
+ self.__ret = self.RETURN_ROTATE_90_COUNTERCLOCKWISE
+ self.__master.destroy()
+
+ def __handle_rotate_90_clockwise(self):
+ self.__ret = self.RETURN_ROTATE_90_CLOCKWISE
+ self.__master.destroy()
+
+ def __handle_next(self):
+ self.__ret = self.RETURN_NEXT
+ self.__master.destroy()
+
+ def __handle_abort(self):
+ self.__ret = self.RETURN_ABORT
+ self.__master.destroy()