Super-squash branch 'main' using huggingface_hub

Co-authored-by: FranckAbgrall <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,34 @@
+FROM python:3.9-slim
+
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    curl \
+    software-properties-common \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+# Try and run pip command after setting the user with `USER user` to avoid permission issues with Python
+RUN pip3 install -r requirements.txt
+
+EXPOSE 8501
+
+HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
+
+ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0", "--server.enableXsrfProtection=false", "--server.enableCORS=false"]

README.md

@@ -0,0 +1,19 @@
+---
+title: Dnbr Tagger Preview1 Demo
+emoji: 🚀
+colorFrom: red
+colorTo: red
+sdk: docker
+app_port: 8501
+tags:
+- streamlit
+pinned: false
+short_description: Demo of a Danbooru tagger
+---
+
+# Welcome to Streamlit!
+
+Edit `/src/streamlit_app.py` to customize this app to your heart's desire. :heart:
+
+If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
+forums](https://discuss.streamlit.io).

requirements.txt

@@ -0,0 +1,6 @@
+altair
+pandas
+streamlit
+timm
+torch
+torchvision

src/streamlit_app.py

@@ -0,0 +1,188 @@
+import streamlit as st
+import numpy as np
+from PIL import Image
+import io  # Used to handle image bytes
+import torch
+import timm
+import json
+from torchvision.transforms.v2 import (
+    ToImage,
+    Compose,
+    ToDtype,
+    Normalize,
+)
+import pandas as pd
+import requests
+
+
+st.set_page_config(layout="wide")
+
+device = ["cpu", "cuda"][torch.cuda.is_available()]
+
+
+def NativeResize(patch_size, seq_len_range):
+    p, lo, hi = patch_size, *seq_len_range
+    refs = sorted(
+        [
+            (i / j, i * p, j * p)
+            for i in range(4, 100)
+            for j in range(4, 100)
+            if 0.33 <= i / j <= 3 and lo <= i * j < hi
+        ]
+    )
+
+    def get_ratio(r):
+        return min(refs, key=lambda rr: max(r, rr[0]) / min(r, rr[0]) - 1)
+
+    def f(im: Image.Image):
+        w, h = im.size
+        _, sw, sh = get_ratio(w / h)
+        return im.resize((sw, sh), resample=Image.Resampling.BICUBIC)
+
+    return f
+
+def load_json_from_url(url):
+    try:
+        response = requests.get(url)
+        response.raise_for_status()  # Raise an exception for bad status codes
+        parsed_json = json.loads(response.text)
+        return parsed_json
+    except requests.exceptions.RequestException as e:
+        print(f"Error fetching data from URL: {e}")
+        return None
+    except json.JSONDecodeError as e:
+        print(f"Error decoding JSON data: {e}")
+        return None
+
+
+@st.cache_data
+def load_tags():
+    freqs = load_json_from_url("https://huggingface.co/gustproof/dnbr-tagger-preview1/raw/main/freqs.json")
+    freqs = [*freqs, (("PLACEHOLDER", 0))]
+    return freqs
+
+
+tags = load_tags()
+
+
+@st.cache_resource
+def load_model():
+    torch.set_grad_enabled(False)
+    model = (
+        timm.create_model(
+            "hf_hub:gustproof/dnbr-tagger-preview1",
+            pretrained=True,
+            dynamic_img_size=True,
+        )
+        .eval()
+        .to(device)
+    )
+    print("loaded model")
+    tf = Compose(
+        [
+            ToImage(),
+            ToDtype(torch.float, scale=True),
+            Normalize(mean=[0.4850, 0.4560, 0.4060], std=[0.2290, 0.2240, 0.2250]),
+        ]
+    )
+
+    class Model:
+        def __init__(self):
+            self.class_names = load_tags()
+
+        def predict(self, img):
+            x = tf(img).unsqueeze(0).to(device)
+            probs = model(x).squeeze(0).sigmoid().cpu()
+            return probs
+
+    return Model()
+
+
+model = load_model()
+
+# --- Streamlit App Layout ---
+st.title("Tagger demo")
+st.write("Model: [gustproof/dnbr-tagger-preview1](https://huggingface.co/gustproof/dnbr-tagger-preview1)")
+st.write(
+    "Upload an image to see predicted labels."
+)
+st.write("---")
+
+# --- Sidebar for Controls ---
+st.sidebar.header("Settings")
+# Confidence Threshold Slider
+confidence_threshold = st.sidebar.slider(
+    "Threshold (recommended: ~0.4-~0.6)",
+    min_value=0.0,
+    max_value=1.0,
+    value=0.5,  # Default threshold
+    step=0.01,
+)
+
+# --- Main Area ---
+col1, col2 = st.columns(2)
+
+with col1:
+    st.header("Upload Image")
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+    if uploaded_file is not None:
+        # Read the image bytes
+        image_bytes = uploaded_file.getvalue()
+        # Display the uploaded image
+        try:
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+            image = NativeResize(14, (270, 301))(image)
+            st.info(f"Resied (width, height): {image.size}")
+            st.image(image, caption="Uploaded Image.")
+        except Exception as e:
+            st.error(f"Error opening image: {e}")
+            st.warning("Please upload a valid image file (JPG, JPEG, PNG).")
+            uploaded_file = None  # Reset uploaded_file so processing stops
+
+with col2:
+    st.header("Predictions")
+    if uploaded_file is not None:
+        with st.spinner("Computing..."):
+            try:
+                scores = model.predict(image)
+                filtered_results = [
+                    (i, p) for i, p in enumerate(scores) if p >= confidence_threshold
+                ]
+                filtered_results.sort(key=lambda x: x[1], reverse=True)
+
+                if filtered_results:
+                    get_category = lambda ti: (
+                        "Rating" if ti < 4 else "General" if ti < 8856 else "Character"
+                    )
+                    df = pd.DataFrame(
+                        [
+                            (i, tags[ti][0], f"{p:.4f}", get_category(ti), tags[ti][1])
+                            for i, (ti, p) in enumerate(filtered_results[:150], 1)
+                        ],
+                        columns=[
+                            "Rank",
+                            "Label",
+                            "Score",
+                            "Category",
+                            "Dataset frequency",
+                        ],
+                    )
+                    st.dataframe(
+                        df,
+                        hide_index=True,
+                        column_config={
+                            "Dataset frequency": st.column_config.NumberColumn(
+                                format="localized"
+                            )
+                        },
+                    )
+                else:
+                    st.info("No labels meet the current threshold.")
+
+            except Exception as e:
+                st.error("An error occurred during prediction or processing:")
+                st.exception(e)  # Shows the full traceback
+
+    else:
+        st.info("Upload an image using the panel on the left to see predictions.")