codeShare/workspace

1from dataclasses import dataclass2from pathlib import Path3from typing import Optional4 5import os6import numpy as np7import pandas as pd8import timm9import torch10from huggingface_hub import hf_hub_download11from huggingface_hub.utils import HfHubHTTPError12from PIL import Image13from simple_parsing import field, parse_known_args14from timm.data import create_transform, resolve_data_config15from torch import Tensor, nn16from torch.nn import functional as F17 18torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")19MODEL_REPO_MAP = {20    "vit": "SmilingWolf/wd-vit-tagger-v3",21    "swinv2": "SmilingWolf/wd-swinv2-tagger-v3",22    "convnext": "SmilingWolf/wd-convnext-tagger-v3",23}24 25 26def pil_ensure_rgb(image: Image.Image) -> Image.Image:27    # convert to RGB/RGBA if not already (deals with palette images etc.)28    if image.mode not in ["RGB", "RGBA"]:29        image = image.convert("RGBA") if "transparency" in image.info else image.convert("RGB")30    # convert RGBA to RGB with white background31    if image.mode == "RGBA":32        canvas = Image.new("RGBA", image.size, (255, 255, 255))33        canvas.alpha_composite(image)34        image = canvas.convert("RGB")35    return image36 37 38def pil_pad_square(image: Image.Image) -> Image.Image:39    w, h = image.size40    # get the largest dimension so we can pad to a square41    px = max(image.size)42    # pad to square with white background43    canvas = Image.new("RGB", (px, px), (255, 255, 255))44    canvas.paste(image, ((px - w) // 2, (px - h) // 2))45    return canvas46 47 48@dataclass49class LabelData:50    names: list[str]51    rating: list[np.int64]52    general: list[np.int64]53    character: list[np.int64]54 55 56def load_labels_hf(57    repo_id: str,58    revision: Optional[str] = None,59    token: Optional[str] = None,60) -> LabelData:61    try:62        csv_path = hf_hub_download(63            repo_id=repo_id, filename="selected_tags.csv", revision=revision, token=token64        )65        csv_path = Path(csv_path).resolve()66    except HfHubHTTPError as e:67        raise FileNotFoundError(f"selected_tags.csv failed to download from {repo_id}") from e68 69    df: pd.DataFrame = pd.read_csv(csv_path, usecols=["name", "category"])70    tag_data = LabelData(71        names=df["name"].tolist(),72        rating=list(np.where(df["category"] == 9)[0]),73        general=list(np.where(df["category"] == 0)[0]),74        character=list(np.where(df["category"] == 4)[0]),75    )76 77    return tag_data78 79 80def get_tags(81    probs: Tensor,82    labels: LabelData,83    gen_threshold: float,84    char_threshold: float,85):86    # Convert indices+probs to labels87    probs = list(zip(labels.names, probs.numpy()))88 89    # First 4 labels are actually ratings90    rating_labels = dict([probs[i] for i in labels.rating])91 92    # General labels, pick any where prediction confidence > threshold93    gen_labels = [probs[i] for i in labels.general]94    gen_labels = dict([x for x in gen_labels if x[1] > gen_threshold])95    gen_labels = dict(sorted(gen_labels.items(), key=lambda item: item[1], reverse=True))96 97    # Character labels, pick any where prediction confidence > threshold98    char_labels = [probs[i] for i in labels.character]99    char_labels = dict([x for x in char_labels if x[1] > char_threshold])100    char_labels = dict(sorted(char_labels.items(), key=lambda item: item[1], reverse=True))101 102    # Combine general and character labels, sort by confidence103    combined_names = [x for x in gen_labels]104    combined_names.extend([x for x in char_labels])105 106    # Convert to a string suitable for use as a training caption107    caption = ", ".join(combined_names)108    taglist = caption.replace("_", " ").replace("(", "\(").replace(")", "\)")109 110    return caption, taglist, rating_labels, char_labels, gen_labels111 112 113@dataclass114class ScriptOptions:115    image_dir: Path = field(positional=True)116    text_dir: Path = field(positional=True)117    model: str = field(default="vit")118    gen_threshold: float = field(default=0.35)119    char_threshold: float = field(default=0.75)120 121 122def main(opts: ScriptOptions):123    repo_id = MODEL_REPO_MAP.get(opts.model)124    print(f"Loading model '{opts.model}' from '{repo_id}'...")125    model: nn.Module = timm.create_model("hf-hub:" + repo_id).eval()    126    state_dict = timm.models.load_state_dict_from_hf(repo_id)127    model.load_state_dict(state_dict)128 129    130    image_dir = Path(opts.image_dir).resolve()131    for image_path in os.findall(image_dir):132        for suffix in ['.jpeg','.jpg','.JPEG','.webp','.WEBP','.png','.PNG']:133            if not image_path.find(suffix)>-1: continue 134            if not image_path.is_file(): raise FileNotFoundError(f"Image file not found: {image_path}")135            136 137 138 139 140 141 142    print("Loading tag list...")143    labels: LabelData = load_labels_hf(repo_id=repo_id)144 145    print("Creating data transform...")146    transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))147 148    print("Loading image and preprocessing...")149    # get image150    img_input: Image.Image = Image.open(image_path)151    # ensure image is RGB152    img_input = pil_ensure_rgb(img_input)153    # pad to square with white background154    img_input = pil_pad_square(img_input)155    # run the model's input transform to convert to tensor and rescale156    inputs: Tensor = transform(img_input).unsqueeze(0)157    # NCHW image RGB to BGR158    inputs = inputs[:, [2, 1, 0]]159 160    print("Running inference...")161    with torch.inference_mode():162        # move model to GPU, if available163        if torch_device.type != "cpu":164            model = model.to(torch_device)165            inputs = inputs.to(torch_device)166        # run the model167        outputs = model.forward(inputs)168        # apply the final activation function (timm doesn't support doing this internally)169        outputs = F.sigmoid(outputs)170        # move inputs, outputs, and model back to to cpu if we were on GPU171        if torch_device.type != "cpu":172            inputs = inputs.to("cpu")173            outputs = outputs.to("cpu")174            model = model.to("cpu")175 176    print("Processing results...")177    caption, taglist, ratings, character, general = get_tags(178        probs=outputs.squeeze(0),179        labels=labels,180        gen_threshold=opts.gen_threshold,181        char_threshold=opts.char_threshold,182    )183 184    print("--------")185    print(f"Caption: {caption}")186    print("--------")187    print(f"Tags: {taglist}")188 189    print("--------")190    print("Ratings:")191    for k, v in ratings.items():192        print(f"  {k}: {v:.3f}")193 194    print("--------")195    print(f"Character tags (threshold={opts.char_threshold}):")196    for k, v in character.items():197        print(f"  {k}: {v:.3f}")198 199    print("--------")200    print(f"General tags (threshold={opts.gen_threshold}):")201    for k, v in general.items():202        print(f"  {k}: {v:.3f}")203 204    print("Done!")205 206 207if __name__ == "__main__":208    opts, _ = parse_known_args(ScriptOptions)209    if opts.model not in MODEL_REPO_MAP:210        print(f"Available models: {list(MODEL_REPO_MAP.keys())}")211        raise ValueError(f"Unknown model name '{opts.model}'")212    main(opts)213