add single file script

scripts/wdtagger3-onnx.py

@@ -0,0 +1,475 @@
+#!/usr/bin/env python3
+import argparse
+import logging
+from dataclasses import dataclass
+from os import PathLike
+from pathlib import Path
+from typing import Generator, Optional, Tuple
+
+import numpy as np
+import onnxruntime as rt
+from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import HfHubHTTPError
+from pandas import DataFrame, read_csv
+from PIL import Image
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+# allowed extensions
+IMAGE_EXTENSIONS = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff", ".tif"]
+# image input shape
+IMAGE_SIZE = 448
+
+MODEL_VARIANTS: dict[str, str] = {
+    "swinv2": "SmilingWolf/wd-swinv2-tagger-v3",
+    "convnext": "SmilingWolf/wd-convnext-tagger-v3",
+    "vit": "SmilingWolf/wd-vit-tagger-v3",
+}
+
+
+@dataclass
+class LabelData:
+    names: list[str]
+    rating: list[np.int64]
+    general: list[np.int64]
+    character: list[np.int64]
+
+
+@dataclass
+class ImageLabels:
+    caption: str
+    booru: str
+    rating: str
+    general: dict[str, float]
+    character: dict[str, float]
+    ratings: dict[str, float]
+
+
+logging.basicConfig(level=logging.WARNING, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+## Model loading functions
+def download_onnx(
+    repo_id: str,
+    filename: str = "model.onnx",
+    revision: Optional[str] = None,
+    token: Optional[str] = None,
+) -> Path:
+    if not filename.endswith(".onnx"):
+        filename += ".onnx"
+
+    model_path = hf_hub_download(repo_id=repo_id, filename=filename, revision=revision, token=token)
+    return Path(model_path).resolve()
+
+
+def create_session(
+    repo_id: str,
+    revision: Optional[str] = None,
+    token: Optional[str] = None,
+) -> rt.InferenceSession:
+    model_path = download_onnx(repo_id, revision=revision, token=token)
+    if not model_path.is_file():
+        model_path = model_path.joinpath("model.onnx")
+    if not model_path.is_file():
+        raise FileNotFoundError(f"Model not found: {model_path}")
+
+    model = rt.InferenceSession(
+        str(model_path),
+        providers=[("CUDAExecutionProvider", {}), "CPUExecutionProvider"],
+    )
+    return model
+
+
+## Label loading function
+def load_labels_hf(
+    repo_id: str,
+    revision: Optional[str] = None,
+    token: Optional[str] = None,
+) -> LabelData:
+    try:
+        csv_path = hf_hub_download(
+            repo_id=repo_id, filename="selected_tags.csv", revision=revision, token=token
+        )
+        csv_path = Path(csv_path).resolve()
+    except HfHubHTTPError as e:
+        raise FileNotFoundError(f"selected_tags.csv failed to download from {repo_id}") from e
+
+    df: DataFrame = read_csv(csv_path, usecols=["name", "category"])
+    tag_data = LabelData(
+        names=df["name"].tolist(),
+        rating=list(np.where(df["category"] == 9)[0]),
+        general=list(np.where(df["category"] == 0)[0]),
+        character=list(np.where(df["category"] == 4)[0]),
+    )
+
+    return tag_data
+
+
+## Image preprocessing functions
+def pil_ensure_rgb(image: Image.Image) -> Image.Image:
+    # convert to RGB/RGBA if not already (deals with palette images etc.)
+    if image.mode not in ["RGB", "RGBA"]:
+        image = image.convert("RGBA") if "transparency" in image.info else image.convert("RGB")
+    # convert RGBA to RGB with white background
+    if image.mode == "RGBA":
+        canvas = Image.new("RGBA", image.size, (255, 255, 255))
+        canvas.alpha_composite(image)
+        image = canvas.convert("RGB")
+    return image
+
+
+def pil_pad_square(
+    image: Image.Image,
+    fill: tuple[int, int, int] = (255, 255, 255),
+) -> Image.Image:
+    w, h = image.size
+    # get the largest dimension so we can pad to a square
+    px = max(image.size)
+    # pad to square with white background
+    canvas = Image.new("RGB", (px, px), fill)
+    canvas.paste(image, ((px - w) // 2, (px - h) // 2))
+    return canvas
+
+
+def preprocess_image(
+    image: Image.Image,
+    size_px: int | tuple[int, int],
+    upscale: bool = True,
+) -> Image.Image:
+    """
+    Preprocess an image to be square and centered on a white background.
+    """
+    if isinstance(size_px, int):
+        size_px = (size_px, size_px)
+
+    # ensure RGB and pad to square
+    image = pil_ensure_rgb(image)
+    image = pil_pad_square(image)
+
+    # resize to target size
+    if image.size[0] < size_px[0] or image.size[1] < size_px[1]:
+        if upscale is False:
+            raise ValueError("Image is smaller than target size, and upscaling is disabled")
+        image = image.resize(size_px, Image.LANCZOS)
+    if image.size[0] > size_px[0] or image.size[1] > size_px[1]:
+        image.thumbnail(size_px, Image.BICUBIC)
+
+    return image
+
+
+## Dataset for DataLoader
+class ImageDataset(Dataset):
+    def __init__(self, image_paths: list[Path], size_px: int = IMAGE_SIZE, upscale: bool = True):
+        self.size_px = size_px
+        self.upscale = upscale
+        self.images = [p for p in image_paths if p.suffix.lower() in IMAGE_EXTENSIONS]
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image_path: Path = self.images[idx]
+        try:
+            image = Image.open(image_path)
+            image = preprocess_image(image, self.size_px, self.upscale)
+            # turn into BGR24 numpy array of N,H,W,C since thats what these want
+            image = image.convert("RGB").convert("BGR;24")
+            image = np.array(image).astype(np.float32)
+        except Exception as e:
+            logging.exception(f"Could not load image from {image_path}, error: {e}")
+            return None
+
+        return {"image": image, "path": np.array(str(image_path).encode("utf-8"), dtype=np.bytes_)}
+
+
+def collate_fn_remove_corrupted(batch):
+    """Collate function that allows to remove corrupted examples in the
+    dataloader. It expects that the dataloader returns 'None' when that occurs.
+    The 'None's in the batch are removed.
+    """
+    # Filter out all the Nones (corrupted examples)
+    batch = [x for x in batch if x is not None]
+    if len(batch) == 0:
+        return None
+    return {k: np.array([x[k] for x in batch if x is not None]) for k in batch[0]}
+
+
+## Main function
+class ImageLabeler:
+    def __init__(
+        self,
+        repo_id: Optional[PathLike] = None,
+        general_threshold: float = 0.35,
+        character_threshold: float = 0.35,
+        banned_tags: list[str] = [],
+    ):
+        self.repo_id = repo_id
+
+        # create some object attributes for convenience
+        self.general_threshold = general_threshold
+        self.character_threshold = character_threshold
+        self.banned_tags = banned_tags if banned_tags is not None else []
+
+        # actually load the model
+        logging.info(f"Loading model from path: {self.repo_id}")
+        self.model = create_session(self.repo_id)
+
+        # Get input dimensions
+        _, self.height, self.width, _ = self.model.get_inputs()[0].shape
+        logging.info(f"Model loaded, input dimensions {self.height}x{self.width}")
+
+        # load labels
+        self.labels = load_labels_hf(self.repo_id)
+        self.labels.general = [i for i in self.labels.general if i not in banned_tags]
+        self.labels.character = [i for i in self.labels.character if i not in banned_tags]
+        logging.info(f"Loaded labels from {self.repo_id}")
+
+    @property
+    def input_size(self) -> Tuple[int, int]:
+        return (self.height, self.width)
+
+    @property
+    def input_name(self) -> str:
+        return self.model.get_inputs()[0].name if self.model is not None else None
+
+    @property
+    def output_name(self) -> str:
+        return self.model.get_outputs()[0].name if self.model is not None else None
+
+    def label_images(self, images: np.ndarray) -> ImageLabels:
+        # Run the ONNX model
+        probs: np.ndarray = self.model.run([self.output_name], {self.input_name: images})[0]
+
+        # Convert to labels
+        results = []
+        for sample in list(probs):
+            labels = list(zip(self.labels.names, sample.astype(float)))
+
+            # First 4 labels are actually ratings: pick one with argmax
+            rating_labels = dict([labels[i] for i in self.labels.rating])
+            rating = max(rating_labels, key=rating_labels.get)
+
+            # General labels, pick any where prediction confidence > threshold
+            gen_labels = [labels[i] for i in self.labels.general]
+            gen_labels = dict([x for x in gen_labels if x[1] > self.general_threshold])
+            gen_labels = dict(sorted(gen_labels.items(), key=lambda item: item[1], reverse=True))
+
+            # Character labels, pick any where prediction confidence > threshold
+            char_labels = [labels[i] for i in self.labels.character]
+            char_labels = dict([x for x in char_labels if x[1] > self.character_threshold])
+            char_labels = dict(sorted(char_labels.items(), key=lambda item: item[1], reverse=True))
+
+            # Combine general and character labels, sort by confidence
+            combined_names = [x for x in gen_labels]
+            combined_names.extend([x for x in char_labels])
+
+            # Convert to a string suitable for use as a training caption
+            caption = ", ".join(combined_names)
+            booru = caption.replace("_", " ").replace("(", "\(").replace(")", "\)")
+
+            # return output
+            results.append(
+                ImageLabels(
+                    caption=caption,
+                    booru=booru,
+                    rating=rating,
+                    general=gen_labels,
+                    character=char_labels,
+                    ratings=rating_labels,
+                )
+            )
+
+        return results
+
+    def __call__(self, images: list[Image.Image]) -> Generator[ImageLabels, None, None]:
+        for x in images:
+            yield self.label_images(x)
+
+
+def main(args):
+    images_dir: Path = Path(args.images_dir).resolve()
+    if not images_dir.is_dir():
+        raise FileNotFoundError(f"Directory not found: {images_dir}")
+
+    variant: str = args.variant
+    recursive: bool = args.recursive or False
+    banned_tags: set[str] = set(args.banned_tags.split(","))
+    caption_extension: str = str(args.caption_extension).lower()
+    print_freqs: bool = args.print_freqs or False
+    num_workers: int = args.num_workers
+    batch_size: int = args.batch_size
+
+    remove_underscore: bool = args.remove_underscore or False
+    general_threshold: float = args.general_threshold or args.thresh
+    character_threshold: float = args.character_threshold or args.thresh
+    debug: bool = args.debug or False
+
+    # turn base model into a repo id and model path
+    repo_id: str = MODEL_VARIANTS.get(variant, None)
+    if repo_id is None:
+        raise ValueError(f"Unknown base model '{variant}'")
+
+    # instantiate the dataset
+    print(f"Loading images from {images_dir}...", end=" ")
+    if recursive is True:
+        image_paths = [p for p in images_dir.rglob("**/*") if p.suffix.lower() in IMAGE_EXTENSIONS]
+    else:
+        image_paths = [p for p in images_dir.glob("*") if p.suffix.lower() in IMAGE_EXTENSIONS]
+
+    n_images = len(image_paths)
+    print(f"found {n_images} images to process, creating DataLoader...")
+    # sort by filename if we have a small number of images
+    if n_images < 10000:
+        image_paths = sorted(image_paths, key=lambda x: x.stem)
+    dataset = ImageDataset(image_paths)
+
+    # Create the data loader
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collate_fn_remove_corrupted,
+        drop_last=False,
+        prefetch_factor=3,
+    )
+
+    # Create the image labeler
+    labeler: ImageLabeler = ImageLabeler(
+        repo_id=repo_id,
+        character_threshold=character_threshold,
+        general_threshold=general_threshold,
+        banned_tags=banned_tags,
+    )
+
+    # object to save tag frequencies
+    tag_freqs = {}
+
+    # iterate
+    for batch in tqdm(dataloader, ncols=100, unit="image", unit_scale=batch_size):
+        images = batch["image"]
+        paths = batch["path"]
+
+        # label the images
+        batch_labels = labeler.label_images(images)
+
+        # save the labels
+        for image_labels, image_path in zip(batch_labels, paths):
+            if isinstance(image_path, (np.bytes_, bytes)):
+                image_path = Path(image_path.decode("utf-8"))
+
+            # save the labels
+            caption = image_labels.caption
+            if remove_underscore is True:
+                caption = caption.replace("_", " ")
+            Path(image_path).with_suffix(caption_extension).write_text(caption + "\n", encoding="utf-8")
+
+            # save the tag frequencies
+            if print_freqs is True:
+                for tag in caption.split(", "):
+                    if tag in banned_tags:
+                        continue
+                    if tag not in tag_freqs:
+                        tag_freqs[tag] = 0
+                    tag_freqs[tag] += 1
+
+            # debug
+            if debug is True:
+                print(
+                    f"{image_path}:"
+                    + f"\n  Character tags: {image_labels.character}"
+                    + f"\n    General tags: {image_labels.general}"
+                )
+
+    if print_freqs:
+        sorted_tags = sorted(tag_freqs.items(), key=lambda x: x[1], reverse=True)
+        print("\nTag frequencies:")
+        for tag, freq in sorted_tags:
+            print(f"{tag}: {freq}")
+
+    print("done!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "images_dir",
+        type=str,
+        help="directory to tag image files in",
+    )
+    parser.add_argument(
+        "--variant",
+        type=str,
+        default="swinv2",
+        help="name of base model to use (one of 'swinv2', 'convnext', 'vit')",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=4,
+        help="number of threads to use in Torch DataLoader (4 should be plenty)",
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=1,
+        help="batch size for Torch DataLoader (use 1 for cpu, 4-32 for gpu)",
+    )
+    parser.add_argument(
+        "--caption_extension",
+        type=str,
+        default=".txt",
+        help="extension of caption files to write (e.g. '.txt', '.caption')",
+    )
+    parser.add_argument(
+        "--thresh",
+        type=float,
+        default=0.35,
+        help="confidence threshold for adding tags",
+    )
+    parser.add_argument(
+        "--general_threshold",
+        type=float,
+        default=None,
+        help="confidence threshold for general tags - defaults to --thresh",
+    )
+    parser.add_argument(
+        "--character_threshold",
+        type=float,
+        default=None,
+        help="confidence threshold for character tags - defaults to --thresh",
+    )
+    parser.add_argument(
+        "--recursive",
+        action="store_true",
+        help="whether to recurse into subdirectories of images_dir",
+    )
+    parser.add_argument(
+        "--remove_underscore",
+        action="store_true",
+        help="whether to remove underscores from tags (e.g. 'long_hair' -> 'long hair')",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="enable debug logging mode",
+    )
+    parser.add_argument(
+        "--banned_tags",
+        type=str,
+        default="",
+        help="tags to filter out (comma-separated)",
+    )
+    parser.add_argument(
+        "--print_freqs",
+        action="store_true",
+        help="Print overall tag frequencies at the end",
+    )
+
+    args = parser.parse_args()
+    if args.images_dir is None:
+        args.images_dir = Path.cwd().joinpath("temp/test")
+
+    main(args)