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)