Update app.py

app.py

@@ -1,319 +1,129 @@
-import argparse
-import os
-import gradio as gr
-import huggingface_hub
+import onnxruntime as ort
 import numpy as np
-import onnxruntime as rt
-import pandas as pd
+import json
 from PIL import Image
-import json  # Added for loading metadata.json from the inference file
 
-TITLE = "WaifuDiffusion Tagger"
-DESCRIPTION = """
-Demo for the WaifuDiffusion tagger models
-
-Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
-"""
-
-HF_TOKEN = os.environ.get("HF_TOKEN", "")
-
-# Dataset v3 series of models:
-SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
-CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
-VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
-VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
-EVA02_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
-
-# Dataset v2 series of models:
-MOAT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
-SWIN_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
-CONV_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
-CONV2_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
-VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
-
-# IdolSankaku series of models:
-EVA02_LARGE_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-eva02-large-tagger-v1"
-SWINV2_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-swinv2-tagger-v1"
-
-# Files to download from the repos
-MODEL_FILENAME = "model.onnx"
-LABEL_FILENAME = "selected_tags.csv"
-
-kaomojis = [
-    "0_0",
-    "(o)_(o)",
-    "+_+",
-    "+_-",
-    "._.",
-    "<o>_<o>",
-    "<|>_<|>",
-    "=_=",
-    ">_<",
-    "3_3",
-    "6_9",
-    ">_o",
-    "@_@",
-    "^_^",
-    "o_o",
-    "u_u",
-    "x_x",
-    "|_|",
-    "||_||",
-]
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--score-slider-step", type=float, default=0.05)
-    parser.add_argument("--score-general-threshold", type=float, default=0.35)
-    parser.add_argument("--score-character-threshold", type=float, default=0.85)
-    return parser.parse_args()
-
-def load_labels(dataframe) -> list[str]:
-    name_series = dataframe["name"]
-    name_series = name_series.map(
-        lambda x: x.replace("_", " ") if x not in kaomojis else x
-    )
-    tag_names = name_series.tolist()
-
-    rating_indexes = list(np.where(dataframe["category"] == 9)[0])
-    general_indexes = list(np.where(dataframe["category"] == 0)[0])
-    character_indexes = list(np.where(dataframe["category"] == 4)[0])
-    return tag_names, rating_indexes, general_indexes, character_indexes
-
-def mcut_threshold(probs):
-    sorted_probs = probs[probs.argsort()[::-1]]
-    difs = sorted_probs[:-1] - sorted_probs[1:]
-    t = difs.argmax()
-    thresh = (sorted_probs[t] + sorted_probs[t + 1]) / 2
-    return thresh
-
-class Predictor:
-    def __init__(self):
-        self.model_target_size = None
-        self.last_loaded_repo = None
-        # Added flag to distinguish between custom and Hugging Face models
-        self.is_custom_model = False
-
-    def download_model(self, model_repo):
-        csv_path = huggingface_hub.hf_hub_download(
-            model_repo,
-            LABEL_FILENAME,
-            use_auth_token=HF_TOKEN,
-        )
-        model_path = huggingface_hub.hf_hub_download(
-            model_repo,
-            MODEL_FILENAME,
-            use_auth_token=HF_TOKEN,
-        )
-        return csv_path, model_path
-
-    def load_model(self, model_repo, onnx_path=None, metadata_path=None):
-        # Modified to accept onnx_path and metadata_path for custom model support
-        if model_repo == "Custom Model" and onnx_path and metadata_path:
-            # Check if the custom model files have already been loaded
-            if self.last_loaded_repo == (onnx_path, metadata_path):
-                return
-            self.is_custom_model = True
-            # Load the ONNX model from the provided path (from inference file)
-            self.model = rt.InferenceSession(onnx_path)
-            # Load metadata from metadata.json (from inference file)
-            with open(metadata_path, "r", encoding="utf-8") as f:
-                metadata = json.load(f)
-            self.idx_to_tag = metadata["idx_to_tag"]
-            # Create tag_names list from idx_to_tag dictionary
-            self.tag_names = [self.idx_to_tag[str(i)] for i in range(len(self.idx_to_tag))]
-            # Set target size to 512 for custom model, as per inference file
-            self.model_target_size = 512
-            self.last_loaded_repo = (onnx_path, metadata_path)
+def preprocess_image(img_path, target_size=512, keep_aspect=True):
+    """
+    Load an image from img_path, convert to RGB,
+    and resize/pad to (target_size, target_size).
+    Scales pixel values to [0,1] and returns a (1,3,target_size,target_size) float32 array.
+    """
+    img = Image.open(img_path).convert("RGB")
+    
+    if keep_aspect:
+        # Preserve aspect ratio, pad black
+        w, h = img.size
+        aspect = w / h
+        if aspect > 1:
+            new_w = target_size
+            new_h = int(new_w / aspect)
         else:
-            # Existing logic for Hugging Face models
-            self.is_custom_model = False
-            if self.last_loaded_repo == model_repo:
-                return
-            csv_path, model_path = self.download_model(model_repo)
-            tags_df = pd.read_csv(csv_path)
-            sep_tags = load_labels(tags_df)
-            self.tag_names = sep_tags[0]
-            self.rating_indexes = sep_tags[1]
-            self.general_indexes = sep_tags[2]
-            self.character_indexes = sep_tags[3]
-            self.model = rt.InferenceSession(model_path)
-            _, height, width, _ = self.model.get_inputs()[0].shape
-            self.model_target_size = height
-            self.last_loaded_repo = model_repo
-
-    def prepare_image(self, image):
-        if self.is_custom_model:
-            # Added preprocessing logic from inference file's preprocess_image function
-            # Adapted to take a PIL image instead of a file path
-            target_size = self.model_target_size
-            img = image.convert("RGB")
-            w, h = img.size
-            aspect = w / h
-            if aspect > 1:
-                new_w = target_size
-                new_h = int(new_w / aspect)
-            else:
-                new_h = target_size
-                new_w = int(new_h * aspect)
-            img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
-            background = Image.new("RGB", (target_size, target_size), (0, 0, 0))
-            paste_x = (target_size - new_w) // 2
-            paste_y = (target_size - new_h) // 2
-            background.paste(img, (paste_x, paste_y))
-            arr = np.array(background).astype("float32") / 255.0
-            arr = np.transpose(arr, (2, 0, 1))  # HWC to CHW as per inference file
-            arr = np.expand_dims(arr, axis=0)
-            return arr
-        else:
-            # Existing preprocessing logic for Hugging Face models
-            target_size = self.model_target_size
-            canvas = Image.new("RGBA", image.size, (255, 255, 255))
-            canvas.alpha_composite(image)
-            image = canvas.convert("RGB")
-            image_shape = image.size
-            max_dim = max(image_shape)
-            pad_left = (max_dim - image_shape[0]) // 2
-            pad_top = (max_dim - image_shape[1]) // 2
-            padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
-            padded_image.paste(image, (pad_left, pad_top))
-            if max_dim != target_size:
-                padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC)
-            image_array = np.asarray(padded_image, dtype=np.float32)
-            image_array = image_array[:, :, ::-1]  # RGB to BGR
-            return np.expand_dims(image_array, axis=0)
-
-    def predict(
-        self,
-        image,
-        model_repo,
-        general_thresh,
-        general_mcut_enabled,
-        character_thresh,
-        character_mcut_enabled,
-        onnx_path=None,
-        metadata_path=None,
-    ):
-        # Modified to accept onnx_path and metadata_path for custom model
-        self.load_model(model_repo, onnx_path, metadata_path)
-        # Added check to ensure custom model files are provided
-        if self.is_custom_model and (onnx_path is None or metadata_path is None):
-            return "Please upload ONNX model and metadata JSON files.", {}, {}, {}
-        image_tensor = self.prepare_image(image)
-        input_name = self.model.get_inputs()[0].name
-        # Changed to use None for output names to get all outputs, supporting custom model
-        outputs = self.model.run(None, {input_name: image_tensor})
-        if self.is_custom_model:
-            # Added inference logic from inference file for custom model
-            # Handle case where model might output initial and refined predictions
-            refined_preds = outputs[1] if len(outputs) == 2 else outputs[0]
-            ref_logit = refined_preds[0]  # Shape (N_tags,)
-            # Apply sigmoid to convert logits to probabilities (from inference file)
-            ref_prob = 1.0 / (1.0 + np.exp(-ref_logit))
-            pred_indices = np.where(ref_prob >= general_thresh)[0]
-            predicted_tags = [self.tag_names[idx] for idx in pred_indices]
-            sorted_general_strings = ", ".join(predicted_tags)
-            # Custom model doesn't use category separation, so return empty for rating and character
-            rating = {}
-            character_res = {}
-            general_res = {self.tag_names[idx]: ref_prob[idx] for idx in pred_indices}
-        else:
-            # Existing inference logic for Hugging Face models
-            preds = outputs[0]  # Assumes single output tensor
-            labels = list(zip(self.tag_names, preds[0].astype(float)))
-            ratings_names = [labels[i] for i in self.rating_indexes]
-            rating = dict(ratings_names)
-            general_names = [labels[i] for i in self.general_indexes]
-            if general_mcut_enabled:
-                general_probs = np.array([x[1] for x in general_names])
-                general_thresh = mcut_threshold(general_probs)
-            general_res = [x for x in general_names if x[1] > general_thresh]
-            general_res = dict(general_res)
-            character_names = [labels[i] for i in self.character_indexes]
-            if character_mcut_enabled:
-                character_probs = np.array([x[1] for x in character_names])
-                character_thresh = mcut_threshold(character_probs)
-                character_thresh = max(0.15, character_thresh)
-            character_res = [x for x in character_names if x[1] > character_thresh]
-            character_res = dict(character_res)
-            sorted_general_strings = sorted(
-                general_res.items(),
-                key=lambda x: x[1],
-                reverse=True,
-            )
-            sorted_general_strings = [x[0] for x in sorted_general_strings]
-            sorted_general_strings = ", ".join(sorted_general_strings).replace("(", r"\(").replace(")", r"\)")
-        return sorted_general_strings, rating, character_res, general_res
-
-def main():
-    args = parse_args()
-    predictor = Predictor()
-    # Added "Custom Model" to the dropdown list to support local ONNX model
-    dropdown_list = [
-        SWINV2_MODEL_DSV3_REPO,
-        CONV_MODEL_DSV3_REPO,
-        VIT_MODEL_DSV3_REPO,
-        VIT_LARGE_MODEL_DSV3_REPO,
-        EVA02_LARGE_MODEL_DSV3_REPO,
-        MOAT_MODEL_DSV2_REPO,
-        SWIN_MODEL_DSV2_REPO,
-        CONV_MODEL_DSV2_REPO,
-        CONV2_MODEL_DSV2_REPO,
-        VIT_MODEL_DSV2_REPO,
-        SWINV2_MODEL_IS_DSV1_REPO,
-        EVA02_LARGE_MODEL_IS_DSV1_REPO,
-        "Custom Model",
-    ]
-    with gr.Blocks(title=TITLE) as demo:
-        with gr.Column():
-            gr.Markdown(value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
-            gr.Markdown(value=DESCRIPTION)
-            with gr.Row():
-                with gr.Column(variant="panel"):
-                    image = gr.Image(type="pil", image_mode="RGBA", label="Input")
-                    model_repo = gr.Dropdown(dropdown_list, value=SWINV2_MODEL_DSV3_REPO, label="Model")
-                    # Added file inputs for ONNX model and metadata, hidden by default
-                    with gr.Row(visible=False) as custom_model_inputs:
-                        onnx_file = gr.File(label="ONNX Model File", file_types=[".onnx"])
-                        metadata_file = gr.File(label="Metadata JSON File", file_types=[".json"])
-                    with gr.Row():
-                        general_thresh = gr.Slider(0, 1, step=args.score_slider_step, value=args.score_general_threshold, label="General Tags Threshold", scale=3)
-                        general_mcut_enabled = gr.Checkbox(value=False, label="Use MCut threshold", scale=1)
-                    with gr.Row():
-                        character_thresh = gr.Slider(0, 1, step=args.score_slider_step, value=args.score_character_threshold, label="Character Tags Threshold", scale=3)
-                        character_mcut_enabled = gr.Checkbox(value=False, label="Use MCut threshold", scale=1)
-                    with gr.Row():
-                        # Updated clear button to include new file inputs
-                        clear = gr.ClearButton(
-                            components=[image, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled, onnx_file, metadata_file],
-                            variant="secondary",
-                            size="lg"
-                        )
-                        submit = gr.Button(value="Submit", variant="primary", size="lg")
-                with gr.Column(variant="panel"):
-                    sorted_general_strings = gr.Textbox(label="Output (string)")
-                    rating = gr.Label(label="Rating")
-                    character_res = gr.Label(label="Output (characters)")
-                    general_res = gr.Label(label="Output (tags)")
-                    clear.add([sorted_general_strings, rating, character_res, general_res])
-        # Added event listener to show/hide custom model inputs based on model selection
-        model_repo.change(
-            lambda x: gr.update(visible=(x == "Custom Model")),
-            inputs=model_repo,
-            outputs=custom_model_inputs,
-        )
-        # Updated submit event to pass onnx_file and metadata_file to predict
-        submit.click(
-            predictor.predict,
-            inputs=[image, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled, onnx_file, metadata_file],
-            outputs=[sorted_general_strings, rating, character_res, general_res],
-        )
-        gr.Examples(
-            [["power.jpg", SWINV2_MODEL_DSV3_REPO, 0.35, False, 0.85, False]],
-            inputs=[image, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled],
-        )
-    demo.queue(max_size=10)
-    demo.launch()
+            new_h = target_size
+            new_w = int(new_h * aspect)
+
+        # Resize with Lanczos
+        img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
+        # Pad to a square
+        background = Image.new("RGB", (target_size, target_size), (0, 0, 0))
+        paste_x = (target_size - new_w) // 2
+        paste_y = (target_size - new_h) // 2
+        background.paste(img, (paste_x, paste_y))
+        img = background
+    else:
+        # simple direct resize to 512x512
+        img = img.resize((target_size, target_size), Image.Resampling.LANCZOS)
+    
+    # Convert to numpy array
+    arr = np.array(img).astype("float32") / 255.0  # scale to [0,1]
+    # Transpose from HWC -> CHW
+    arr = np.transpose(arr, (2, 0, 1))
+    # Add batch dimension: (1,3,512,512)
+    arr = np.expand_dims(arr, axis=0)
+    return arr
+
+def onnx_inference(img_paths, 
+                   onnx_path="camie_refined_no_flash.onnx", 
+                   threshold=0.325, 
+                   metadata_file="metadata.json"):
+    """
+    Loads the ONNX model, runs inference on a list of image paths,
+    and applies an optional threshold to produce final predictions.
+    
+    Args:
+      img_paths: List of paths to images.
+      onnx_path: Path to the exported ONNX model file.
+      threshold: Probability threshold for deciding if a tag is predicted.
+      metadata_file: Path to metadata.json that contains idx_to_tag etc.
+    
+    Returns:
+      A list of dicts, each containing:
+        {
+          "initial_logits": np.ndarray of shape (N_tags,),
+          "refined_logits": np.ndarray of shape (N_tags,),
+          "predicted_tags": list of tag indices that exceeded threshold,
+          ...
+        }
+      one dict per input image.
+    """
+    # 1) Initialize ONNX runtime session
+    session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+    # Optional: for GPU usage, see if "CUDAExecutionProvider" is available
+    # session = ort.InferenceSession(onnx_path, providers=["CUDAExecutionProvider"])
+
+    # 2) Pre-load metadata
+    with open(metadata_file, "r", encoding="utf-8") as f:
+        metadata = json.load(f)
+    idx_to_tag = metadata["idx_to_tag"]  # e.g. { "0": "brown_hair", "1": "blue_eyes", ... }
+
+    # 3) Preprocess each image into a batch
+    batch_tensors = []
+    for img_path in img_paths:
+        x = preprocess_image(img_path, target_size=512, keep_aspect=True)
+        batch_tensors.append(x)
+    # Concatenate along the batch dimension => shape (batch_size, 3, 512, 512)
+    batch_input = np.concatenate(batch_tensors, axis=0)
+
+    # 4) Run inference
+    input_name = session.get_inputs()[0].name        # typically "image"
+    outputs = session.run(None, {input_name: batch_input})
+    # Typically we get [initial_tags, refined_tags] as output
+    initial_preds, refined_preds = outputs  # shapes => (batch_size, 70527)
+
+    # 5) For each image in batch, convert logits to predictions if desired
+    batch_results = []
+    for i in range(initial_preds.shape[0]):
+        # Extract one sample's logits
+        init_logit = initial_preds[i, :]   # shape (N_tags,)
+        ref_logit = refined_preds[i, :]    # shape (N_tags,)
+
+        # Convert to probabilities with sigmoid
+        ref_prob = 1.0 / (1.0 + np.exp(-ref_logit))
+
+        # Threshold
+        pred_indices = np.where(ref_prob >= threshold)[0]
+
+        # Build result for this image
+        result_dict = {
+            "initial_logits": init_logit,
+            "refined_logits": ref_logit,
+            "predicted_indices": pred_indices,
+            "predicted_tags": [idx_to_tag[str(idx)] for idx in pred_indices]  # map index->tag name
+        }
+        batch_results.append(result_dict)
+
+    return batch_results
 
 if __name__ == "__main__":
-    main()
+    # Example usage
+    images = ["image1.jpg", "image2.jpg", "image3.jpg"]
+    results = onnx_inference(images,
+                             onnx_path="camie_refined_no_flash.onnx",
+                             threshold=0.325,
+                             metadata_file="metadata.json")
+
+    for i, res in enumerate(results):
+        print(f"Image: {images[i]}")
+        print(f"  # of predicted tags above threshold: {len(res['predicted_indices'])}")
+        print(f"  Some predicted tags: {res['predicted_tags'][:10]}  (Show up to 10)")
+        print()