Initial commit.

.gitignore

@@ -0,0 +1 @@
+__pycache__

Dockerfile

@@ -0,0 +1,31 @@
+FROM python:3.9.13-slim
+WORKDIR /autotagger
+
+# https://github.com/python-poetry/poetry/discussions/1879#discussioncomment-216865
+ENV \
+  # https://stackoverflow.com/questions/59812009/what-is-the-use-of-pythonunbuffered-in-docker-file
+  PYTHONUNBUFFERED=1 \
+  # https://python-docs.readthedocs.io/en/latest/writing/gotchas.html#disabling-bytecode-pyc-files
+  PYTHONDONTWRITEBYTECODE=1 \
+  # https://stackoverflow.com/questions/45594707/what-is-pips-no-cache-dir-good-for
+  PIP_NO_CACHE_DIR=1 \
+  # https://stackoverflow.com/questions/46288847/how-to-suppress-pip-upgrade-warning
+  PIP_DISABLE_PIP_VERSION_CHECK=1 \
+  PATH=/autotagger:$PATH
+
+RUN \
+  apt-get update && \
+  apt-get install -y --no-install-recommends tini build-essential gfortran libatlas-base-dev && \
+  pip install "poetry==1.1.13"
+
+COPY pyproject.toml poetry.lock ./
+RUN \
+  python -m poetry install --no-dev && \
+  rm -rf /root/.cache/pypoetry/artifacts /root/.cache/pypoetry/cache
+COPY . .
+
+EXPOSE 5000
+ENTRYPOINT ["tini", "--", "poetry", "run"]
+#CMD ["autotag"]
+#CMD ["flask", "run", "--host", "0.0.0.0"]
+CMD ["gunicorn", "app:app", "--bind", "0.0.0.0:5000"]

LICENSE

@@ -0,0 +1,17 @@
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+
+from os import getenv
+from dotenv import load_dotenv
+from autotagger import Autotagger
+from base64 import b64encode
+from flask import Flask, request, render_template, jsonify
+
+load_dotenv()
+model_path = getenv("MODEL_PATH", "models/model.pth")
+autotagger = Autotagger(model_path)
+
+app = Flask(__name__)
+app.config["JSON_SORT_KEYS"] = False
+app.config["JSON_PRETTYPRINT_REGULAR"] = True
+
+@app.route("/", methods=["GET"])
+def index():
+    return render_template("index.html")
+
+@app.route("/evaluate", methods=["POST"])
+def evaluate():
+    files = request.files.getlist("file")
+    threshold = float(request.form.get("threshold", 0.1))
+    output = request.form.get("format", "html")
+    limit = int(request.form.get("limit", 50))
+
+    if output == "html":
+        predictions = [{ "data": b64encode(data).decode(), "tags": autotagger.predict(data, threshold=threshold, limit=limit) } for data in (file.stream.read() for file in files)]
+        return render_template("evaluate.html", predictions=predictions)
+    elif output == "json":
+        predictions = [{ "filename": file.filename, "tags": autotagger.predict(file.read(), threshold=threshold, limit=limit) } for file in files]
+        return jsonify(predictions)
+    else:
+        return 400
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0")

autotag

@@ -0,0 +1,19 @@
+#!/usr/bin/env python
+
+import json
+import click
+from autotagger import Autotagger
+from fastai.vision.core import PILImage
+
+@click.command(help="Automatically generate tags for an image.")
+@click.option("-t", "--threshold", default=0.01, type=float, show_default=True, help="The minimum tag confidence level.")
+@click.option("-n", "--limit", default=50, type=int, show_default=True, help="The maximum number of tags to return.")
+@click.option("-m", "--model", default="models/model.pth", type=click.Path(exists=True), help="The model to use.")
+@click.argument("file", nargs=-1, type=click.File("rb"), required=True)
+def main(file, threshold, limit, model):
+    autotagger = Autotagger(model)
+    predictions = [{ "filename": f.name, "tags": autotagger.predict(PILImage.create(f), threshold=threshold, limit=limit) } for f in file]
+    click.echo(json.dumps(predictions, indent=2))
+
+if __name__ == "__main__":
+    main()

autotagger/__init__.py

@@ -0,0 +1 @@
+from .autotagger import Autotagger

autotagger/autotagger.py

@@ -0,0 +1,35 @@
+from fastbook import *
+from pandas import read_csv
+import timm
+
+class Autotagger:
+    def __init__(self, model_path="models/model.pth", data_path="test/tags.csv.gz", tags_path="data/tags.json"):
+        self.model_path = model_path
+        self.learn = self.init_model(data_path=data_path, tags_path=tags_path, model_path=model_path)
+
+    def init_model(self, model_path="model/model.pth", data_path="test/tags.csv.gz", tags_path="data/tags.json"):
+        df = read_csv(data_path)
+        vocab = json.load(open(tags_path))
+
+        dblock = DataBlock(
+            blocks=(ImageBlock, MultiCategoryBlock(vocab=vocab)),
+            get_x = lambda df: Path("test") / df["filename"],
+            get_y = lambda df: df["tags"].split(" "),
+            item_tfms = Resize(224, method = ResizeMethod.Squish),
+            batch_tfms = [RandomErasing()]
+        )
+
+        dls = dblock.dataloaders(df)
+        learn = vision_learner(dls, "resnet152", pretrained=False)
+        model_file = open(model_path, "rb")
+        learn.load(model_file, with_opt=False)
+
+        return learn
+
+    def predict(self, path, threshold=0.01, limit=50):
+        with self.learn.no_bar(), self.learn.no_logging():
+            pred = self.learn.predict(path)
+            scores = [score.item() for score in pred[2]]
+            results = { tag : score for (tag, score) in zip(self.learn.dls.vocab, scores) if score >= threshold }
+            results = sorted(results.items(), key = lambda x: -x[1])
+            return dict(results[:limit])

pyproject.toml

@@ -0,0 +1,23 @@
+[tool.poetry]
+name = "autotagger"
+version = "1.0"
+description = "Danbooru autotagger"
+authors = ["evazion <noizave@gmail.com>"]
+license = "MIT"
+
+[tool.poetry.dependencies]
+python = "~3.9"
+fastbook = "^0.0.26"
+Flask = "^2.1.2"
+python-dotenv = "^0.20.0"
+click = "^8.1.3"
+ipywidgets = "^7.7.0"
+timm = "^0.5.4"
+scipy = "^1.8.1"
+gunicorn = "^20.1.0"
+
+[tool.poetry.dev-dependencies]
+
+[build-system]
+requires = ["poetry>=0.12"]
+build-backend = "poetry.masonry.api"

templates/evaluate.html

@@ -0,0 +1,37 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <script src="https://cdn.tailwindcss.com"></script>
+  </head>
+
+  <body class="text-sm m-4 break-all lg:max-w-[960px] lg:mx-auto" style="font-family: system-ui;">
+    <h1 class="text-3xl">Results</h1>
+    <a class="text-xs text-sky-600 hover:text-sky-500 mr-4" href="/">&lt; Back</a>
+
+    <div class="mt-4">
+      {% for prediction in predictions %}
+        <div class="flex flex-col p-2 gap-2 border rounded md:flex-row md:max-h-[80vh]">
+          <div class="flex-1 flex items-center justify-center">
+            <img class="max-w-full max-h-full h-auto" src="data:image/jpg;base64,{{ prediction["data"] | safe }}">
+          </div>
+
+          <div class="flex-0 overflow-scroll md:pr-2">
+            <table class="w-full leading-4">
+              {% for tag, score in prediction["tags"].items() %}
+              <tr>
+                <td>
+                  <a class="text-sky-600 hover:text-sky-500" href="https://danbooru.donmai.us/wiki_pages/{{ tag | urlencode }}">?</a>
+                  <a class="text-sky-600 hover:text-sky-500 mr-4" href="https://danbooru.donmai.us/posts?tags={{ tag | urlencode }}">{{ tag | replace("_", " ") }}</a>
+                </td>
+                <td class="text-gray-400 text-right">{{ "{:.0f}%".format(100 * score) }}</td>
+              </tr>
+              {% endfor %}
+            </table>
+
+            <textarea class="w-full text-gray-500 mt-2" rows="4">{{ " ".join(prediction["tags"].keys()) }}</textarea>
+          </div>
+        </div>
+      {% endfor %}
+    </div>
+  </body>
+</html>

templates/index.html

@@ -0,0 +1,11 @@
+<!DOCTYPE html>
+<html>
+  <body>
+    <form action="/evaluate" method="post" enctype="multipart/form-data">
+      <input type="file" name="file">
+      <input type="hidden" name="threshold" min="0" max="1" step="0.1" value="0.01">
+      <input type="hidden" name="limit" value="100">
+      <input type="submit" value="Submit">
+    </form>
+  </body>
+</html>