```
from logging import getLogger
from pathlib import Path
from typing import Generator, Iterable
import sqlite3
import json
from threading import Lock

from PIL import Image
from pydantic import BaseModel, Field
from tqdm import tqdm

from configs.path_config import CACHE_DB_DIR, CACHE_DIR
from tagger.interrogator.interrogator import AbsInterrogator
from tagger.interrogators import interrogators
from tagger.tag_merger import merge_and_save_tags

logger_local = getLogger(__name__)

_db_lock = Lock()
def init_cache_db() -> None:
    """
    Initialize the cache database and create table if not exists.
    """
    CACHE_DB_DIR.parent.mkdir(parents=True, exist_ok=True)
    with _db_lock, sqlite3.connect(CACHE_DB_DIR) as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS tag_cache (
                image_id INTEGER PRIMARY KEY,
                tags_json TEXT
            )
            """
        )
        conn.commit()
def get_cached_tags(image_id: int, model_name: str) -> list[str] | None:
    """
    Retrieve cached tags for a given image_id and model_name, if present.
    """
    with _db_lock, sqlite3.connect(CACHE_DB_DIR) as conn:
        cursor = conn.execute(
            "SELECT tags_json FROM tag_cache WHERE image_id = ?",
            (image_id,)
        )
        row = cursor.fetchone()
        if not row:
            return None
        data = json.loads(row[0])  # dict of model_name -> tag string
        tag_str = data.get(model_name)
        if not tag_str:
            return None
        # Split and strip tags
        return [t.strip() for t in tag_str.split(',') if t.strip()]
def set_cached_tags(image_id: int, model_name: str, tags: list[str]) -> None:
    """
    Store tags list for a given image_id and model_name into cache.
    """
    with _db_lock, sqlite3.connect(CACHE_DB_DIR) as conn:
        cursor = conn.execute(
            "SELECT tags_json FROM tag_cache WHERE image_id = ?",
            (image_id,)
        )
        row = cursor.fetchone()
        if row:
            data = json.loads(row[0])
        else:
            data = {}
        data[model_name] = ", ".join(tags)
        json_text = json.dumps(data, ensure_ascii=False)
        if row:
            conn.execute(
                "UPDATE tag_cache SET tags_json = ? WHERE image_id = ?",
                (json_text, image_id)
            )
        else:
            conn.execute(
                "INSERT INTO tag_cache(image_id, tags_json) VALUES(?, ?)"
                , (image_id, json_text)
            )
        conn.commit()
# Initialize the cache DB on import
init_cache_db()
class TagEntry(BaseModel):
    data: dict[str, str] = Field(default={})
    # {
    #   "danbooru": "tag, tag, tag,...",
    #   "model_name": "tag, tag, tag..."
    # }
def parse_exclude_tags(exclude_tags: list[str]) -> set[str]:
    """
    Parse and format exclude tags.
    """
    if exclude_tags is None:
        return set()

    tags = []
    for tag_str in exclude_tags:
        for tag in tag_str.split(','):
            tags.append(tag.strip())

    # Reverse escape (nai tag to danbooru tag)
    reverse_escaped_tags = []
    for tag in tags:
        t = tag.replace(' ', '_').replace(r'\(', '(').replace(r'\)', ')')
        reverse_escaped_tags.append(t)
    return {*tags, *reverse_escaped_tags}
async def image_interrogate(
        image_path: Path,
        interrogator_model,
        threshold: float,
        tag_escape: bool,
        exclude_tags: Iterable[str]
) -> dict[str, float]:
    """
    Predictions from an image path.
    Bypass inference if cached; otherwise, check for same-named image in CACHE_DIR first.
    """
    # Determine image id from filename stem
    try:
        image_id = int(image_path.stem)
    except ValueError:
        image_id = hash(image_path.name)
    model_name = interrogator_model.name

    # Check cache
    cached = get_cached_tags(image_id, model_name)
    if cached is not None:
        return {tag: 1.0 for tag in cached}

    # If not cached, check for same-named image in CACHE_DIR
    cache_image_path = CACHE_DIR / image_path.name
    if cache_image_path.is_file():
        actual_path = cache_image_path
    else:
        actual_path = image_path

    logger_local.info(f"Performing inference on image: {actual_path}")

    # Perform inference
    im = Image.open(actual_path)
    result = await interrogator_model.interrogate(im)
    tags_dict = AbsInterrogator.postprocess_tags(
        result[1],
        threshold=threshold,
        escape_tag=tag_escape,
        replace_underscore=tag_escape,
        exclude_tags=exclude_tags
    )
    tags = list(tags_dict.keys())

    # Cache results
    set_cached_tags(image_id, model_name, tags)

    return tags_dict
def explore_image_files(folder_path: Path, recursive: bool = False) -> Generator[Path, None, None]:
    """
    Explore image files by folder path.
    """
    for path in folder_path.iterdir():
        if path.is_file() and path.suffix in ['.png', '.jpg', '.jpeg', '.webp']:
            yield path
        elif recursive and path.is_dir():
            yield from explore_image_files(path, recursive)
async def process_images_in_directory(
        dir_path: str,
        model: str = 'wd-vit-large-tagger-v3',
        threshold: float = 0.35,
        ext: str = '.txt',
        overwrite: bool = False,
        recursive: bool = False,
        rawtag: bool = False,
        exclude_tags: list[str] = None,
        use_cpu: bool = False
):
    """
    Process all images in a directory and generate captions.
    """
    root_path = Path(dir_path)
    interrogator = interrogators[model]

    # Use CPU if specified
    if use_cpu:
        interrogator.use_cpu()

    # Parse exclude tags
    exclude_tags_set = parse_exclude_tags(exclude_tags)

    # Gather generator
    image_paths = list(explore_image_files(root_path, recursive))
    total_images = len(image_paths)

    # Process each images
    for image_path in tqdm(
            image_paths,
            desc="Processing images",
            total=total_images,
            unit="image",
            dynamic_ncols=True
    ):
        caption_path = image_path.parent / f'{image_path.stem}{ext}'
        logger_local.debug(f"Processing {image_path}")

        tags = await image_interrogate(
            image_path,
            interrogator,
            threshold,
            not rawtag,
            exclude_tags_set
        )
        detected_tags = set(tags.keys())

        if caption_path.is_file() and not overwrite:
            merge_and_save_tags(caption_path, detected_tags)
        else:
            tags_str = ', '.join(detected_tags)
            with open(caption_path, 'w', encoding='utf-8') as fp:
                fp.write(tags_str)
async def process_single_image(
        file_path: str,
        model: str = 'wd-vit-large-tagger-v3',
        threshold: float = 0.35,
        ext: str = '.txt',
        overwrite: bool = False,
        rawtag: bool = False,
        exclude_tags: list[str] = None,
        use_cpu: bool = False
) -> str:
    """
    Process a single image and generate a caption.
    """
    image_path = Path(file_path)
    interrogator = interrogators[model]

    # Use CPU if specified
    if use_cpu:
        interrogator.use_cpu()

    # Parse exclude tags
    exclude_tags_set = parse_exclude_tags(exclude_tags)

    # Interrogate and log path
    tags_dict = await image_interrogate(
        image_path,
        interrogator,
        threshold,
        not rawtag,
        exclude_tags_set
    )
    tags_str = ', '.join(tags_dict.keys())

    # Save tags to file
    caption_path = image_path.with_suffix(ext)
    if caption_path.is_file() and not overwrite:
        merge_and_save_tags(caption_path, set(tags_dict.keys()))
    else:
        with open(caption_path, 'w', encoding='utf-8') as fp:
            fp.write(tags_str)

    return tags_str
```
データベースに関する処理を別ファイルに分けて、それぞれの新しいコード全体を書いて。