2021-04-15 14:21:36 +03:00
|
|
|
import argparse
|
2021-04-15 12:30:26 +03:00
|
|
|
import logging
|
|
|
|
import os
|
|
|
|
import re
|
|
|
|
import shutil
|
|
|
|
import time
|
|
|
|
from datetime import datetime
|
|
|
|
from pathlib import Path
|
2021-04-15 14:21:36 +03:00
|
|
|
from typing import Optional
|
2021-04-15 12:30:26 +03:00
|
|
|
|
|
|
|
from config import Config
|
|
|
|
from database import Database
|
2021-04-14 23:23:56 +03:00
|
|
|
from iqdb import Iqdb
|
|
|
|
from library import Library
|
|
|
|
from metadata import Metadata
|
|
|
|
|
|
|
|
|
|
|
|
class PicSorter:
|
2021-04-15 14:21:36 +03:00
|
|
|
@staticmethod
|
|
|
|
def parse_args():
|
|
|
|
parser = argparse.ArgumentParser(
|
|
|
|
description='Finds an image on danbooru, writes tags as IPTC keywords, than places the image in the library'
|
|
|
|
)
|
|
|
|
parser.add_argument('-c', '--config',
|
|
|
|
type=Path,
|
|
|
|
default='config.yml',
|
|
|
|
help='config.yml file path')
|
|
|
|
parser.add_argument('input', nargs=argparse.REMAINDER)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if len(args.input) >= 1:
|
|
|
|
PicSorter(args.config).process(args.input)
|
|
|
|
|
|
|
|
def __init__(self, config_file='config.yml'):
|
|
|
|
config = Config.load(config_file)
|
|
|
|
self.config = config
|
|
|
|
self.__setup_logging(config.dir_logs)
|
|
|
|
self.library = Library(config.dir_library)
|
|
|
|
self.metadata = Metadata(config.dir_tmp)
|
|
|
|
self.db = Database()
|
2021-04-14 23:23:56 +03:00
|
|
|
|
2021-04-15 14:21:36 +03:00
|
|
|
@staticmethod
|
|
|
|
def __setup_logging(dir_logs: Path):
|
2021-04-15 12:30:26 +03:00
|
|
|
filename = datetime.now().strftime('%Y-%m-%d.log')
|
2021-04-15 14:21:36 +03:00
|
|
|
logfile = Path(dir_logs, filename)
|
2021-04-14 23:23:56 +03:00
|
|
|
logging.basicConfig(
|
2021-04-15 12:30:26 +03:00
|
|
|
filename=os.fspath(logfile),
|
2021-04-14 23:23:56 +03:00
|
|
|
level=logging.INFO,
|
|
|
|
format='%(asctime)s %(levelname)s %(module)s: %(message)s',
|
|
|
|
datefmt='%H:%M:%S',
|
|
|
|
)
|
|
|
|
|
2021-04-15 14:21:36 +03:00
|
|
|
def process(self, inputs: list[str]) -> None:
|
|
|
|
for input in inputs:
|
2021-04-18 13:44:50 +03:00
|
|
|
if input.startswith("http"):
|
2021-04-15 14:21:36 +03:00
|
|
|
print("Processing url", input)
|
|
|
|
self.__process_url(input)
|
2021-04-18 13:44:50 +03:00
|
|
|
else:
|
|
|
|
p = Path(input)
|
|
|
|
if p.is_dir():
|
|
|
|
self.__process_folder(p)
|
|
|
|
elif p.is_file():
|
|
|
|
print("Processing file", input)
|
|
|
|
self.__process_file(input)
|
2021-04-15 14:21:36 +03:00
|
|
|
|
|
|
|
def __process_folder(self, dir_input: Path) -> None:
|
|
|
|
files = {p for p in dir_input.iterdir()
|
2021-04-15 12:30:26 +03:00
|
|
|
if p.suffix in [".jpg", ".png"]}
|
2021-04-15 00:38:10 +03:00
|
|
|
for filename in files:
|
2021-04-15 14:21:36 +03:00
|
|
|
print("Processing", filename)
|
2021-04-14 23:23:56 +03:00
|
|
|
try:
|
2021-04-15 14:21:36 +03:00
|
|
|
self.__process_file(filename)
|
2021-04-14 23:23:56 +03:00
|
|
|
except Exception as ex:
|
|
|
|
raise ex
|
2021-04-15 14:21:36 +03:00
|
|
|
time.sleep(5)
|
|
|
|
|
|
|
|
def __process_file(self, filename: str) -> bool:
|
|
|
|
url = self.__search_iqdb(filename)
|
|
|
|
if url is None:
|
|
|
|
return False
|
|
|
|
if self.__process_url(url):
|
|
|
|
self.config.dir_processed.mkdir(exist_ok=True, parents=True)
|
|
|
|
from_path = os.fspath(filename)
|
|
|
|
to_path = os.fspath(self.config.dir_processed)
|
|
|
|
shutil.move(from_path, to_path)
|
|
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
|
|
def __search_iqdb(self, filename: str) -> Optional[str]:
|
|
|
|
url = Iqdb.search(filename)
|
|
|
|
if url is None:
|
|
|
|
logging.warning("%s not found", filename)
|
|
|
|
self.library.move_to_orphan(Path(filename))
|
|
|
|
return None
|
|
|
|
return url
|
|
|
|
|
|
|
|
def __process_url(self, url: str) -> bool:
|
|
|
|
m = re.search(r".*posts/(\d{3,})", url)
|
|
|
|
if not m:
|
|
|
|
return False
|
|
|
|
post_id = int(m.group(1))
|
|
|
|
if self.db.is_exists(post_id):
|
|
|
|
logging.info("Skipping exists post %d", post_id)
|
|
|
|
return False
|
|
|
|
|
|
|
|
meta_result = self.metadata.process(url)
|
|
|
|
if meta_result is None:
|
|
|
|
return False
|
|
|
|
image_path, tags = meta_result
|
|
|
|
self.library.move(image_path, tags)
|
|
|
|
self.db.add(post_id, tags.tags_string)
|
|
|
|
return True
|
2021-04-15 12:30:26 +03:00
|
|
|
|
|
|
|
|
2021-04-14 23:23:56 +03:00
|
|
|
if __name__ == '__main__':
|
2021-04-15 14:21:36 +03:00
|
|
|
PicSorter.parse_args()
|