init scan

.gitignore

@@ -0,0 +1,479 @@
+## Ignore Visual Studio temporary files, build results, and
+## files generated by popular Visual Studio add-ons.
+##
+## Get latest from https://github.com/github/gitignore/blob/main/VisualStudio.gitignore
+
+# User-specific files
+*.rsuser
+*.suo
+*.user
+*.userosscache
+*.sln.docstates
+
+# User-specific files (MonoDevelop/Xamarin Studio)
+*.userprefs
+
+# Mono auto generated files
+mono_crash.*
+
+# Build results
+[Dd]ebug/
+[Dd]ebugPublic/
+[Rr]elease/
+[Rr]eleases/
+x64/
+x86/
+[Ww][Ii][Nn]32/
+[Aa][Rr][Mm]/
+[Aa][Rr][Mm]64/
+bld/
+[Bb]in/
+[Oo]bj/
+[Ll]og/
+[Ll]ogs/
+
+# Visual Studio 2015/2017 cache/options directory
+.vs/
+# Uncomment if you have tasks that create the project's static files in wwwroot
+#wwwroot/
+
+# Visual Studio 2017 auto generated files
+Generated\ Files/
+
+# MSTest test Results
+[Tt]est[Rr]esult*/
+[Bb]uild[Ll]og.*
+
+# NUnit
+*.VisualState.xml
+TestResult.xml
+nunit-*.xml
+
+# Build Results of an ATL Project
+[Dd]ebugPS/
+[Rr]eleasePS/
+dlldata.c
+
+# Benchmark Results
+BenchmarkDotNet.Artifacts/
+
+# .NET
+project.lock.json
+project.fragment.lock.json
+artifacts/
+
+# Tye
+.tye/
+
+# ASP.NET Scaffolding
+ScaffoldingReadMe.txt
+
+# StyleCop
+StyleCopReport.xml
+
+# Files built by Visual Studio
+*_i.c
+*_p.c
+*_h.h
+*.ilk
+*.meta
+*.obj
+*.iobj
+*.pch
+*.pdb
+*.ipdb
+*.pgc
+*.pgd
+*.rsp
+*.sbr
+*.tlb
+*.tli
+*.tlh
+*.tmp
+*.tmp_proj
+*_wpftmp.csproj
+*.log
+*.tlog
+*.vspscc
+*.vssscc
+.builds
+*.pidb
+*.svclog
+*.scc
+
+# Chutzpah Test files
+_Chutzpah*
+
+# Visual C++ cache files
+ipch/
+*.aps
+*.ncb
+*.opendb
+*.opensdf
+*.sdf
+*.cachefile
+*.VC.db
+*.VC.VC.opendb
+
+# Visual Studio profiler
+*.psess
+*.vsp
+*.vspx
+*.sap
+
+# Visual Studio Trace Files
+*.e2e
+
+# TFS 2012 Local Workspace
+$tf/
+
+# Guidance Automation Toolkit
+*.gpState
+
+# ReSharper is a .NET coding add-in
+_ReSharper*/
+*.[Rr]e[Ss]harper
+*.DotSettings.user
+
+# TeamCity is a build add-in
+_TeamCity*
+
+# DotCover is a Code Coverage Tool
+*.dotCover
+
+# AxoCover is a Code Coverage Tool
+.axoCover/*
+!.axoCover/settings.json
+
+# Coverlet is a free, cross platform Code Coverage Tool
+coverage*.json
+coverage*.xml
+coverage*.info
+
+# Visual Studio code coverage results
+*.coverage
+*.coveragexml
+
+# NCrunch
+_NCrunch_*
+.*crunch*.local.xml
+nCrunchTemp_*
+
+# MightyMoose
+*.mm.*
+AutoTest.Net/
+
+# Web workbench (sass)
+.sass-cache/
+
+# Installshield output folder
+[Ee]xpress/
+
+# DocProject is a documentation generator add-in
+DocProject/buildhelp/
+DocProject/Help/*.HxT
+DocProject/Help/*.HxC
+DocProject/Help/*.hhc
+DocProject/Help/*.hhk
+DocProject/Help/*.hhp
+DocProject/Help/Html2
+DocProject/Help/html
+
+# Click-Once directory
+publish/
+
+# Publish Web Output
+*.[Pp]ublish.xml
+*.azurePubxml
+# Note: Comment the next line if you want to checkin your web deploy settings,
+# but database connection strings (with potential passwords) will be unencrypted
+*.pubxml
+*.publishproj
+
+# Microsoft Azure Web App publish settings. Comment the next line if you want to
+# checkin your Azure Web App publish settings, but sensitive information contained
+# in these scripts will be unencrypted
+PublishScripts/
+
+# NuGet Packages
+*.nupkg
+# NuGet Symbol Packages
+*.snupkg
+# The packages folder can be ignored because of Package Restore
+**/[Pp]ackages/*
+# except build/, which is used as an MSBuild target.
+!**/[Pp]ackages/build/
+# Uncomment if necessary however generally it will be regenerated when needed
+#!**/[Pp]ackages/repositories.config
+# NuGet v3's project.json files produces more ignorable files
+*.nuget.props
+*.nuget.targets
+
+# Microsoft Azure Build Output
+csx/
+*.build.csdef
+
+# Microsoft Azure Emulator
+ecf/
+rcf/
+
+# Windows Store app package directories and files
+AppPackages/
+BundleArtifacts/
+Package.StoreAssociation.xml
+_pkginfo.txt
+*.appx
+*.appxbundle
+*.appxupload
+
+# Visual Studio cache files
+# files ending in .cache can be ignored
+*.[Cc]ache
+# but keep track of directories ending in .cache
+!?*.[Cc]ache/
+
+# Others
+ClientBin/
+~$*
+*~
+*.dbmdl
+*.dbproj.schemaview
+*.jfm
+*.pfx
+*.publishsettings
+orleans.codegen.cs
+
+# Including strong name files can present a security risk
+# (https://github.com/github/gitignore/pull/2483#issue-259490424)
+#*.snk
+
+# Since there are multiple workflows, uncomment next line to ignore bower_components
+# (https://github.com/github/gitignore/pull/1529#issuecomment-104372622)
+#bower_components/
+
+# RIA/Silverlight projects
+Generated_Code/
+
+# Backup & report files from converting an old project file
+# to a newer Visual Studio version. Backup files are not needed,
+# because we have git ;-)
+_UpgradeReport_Files/
+Backup*/
+UpgradeLog*.XML
+UpgradeLog*.htm
+ServiceFabricBackup/
+*.rptproj.bak
+
+# SQL Server files
+*.mdf
+*.ldf
+*.ndf
+
+# Business Intelligence projects
+*.rdl.data
+*.bim.layout
+*.bim_*.settings
+*.rptproj.rsuser
+*- [Bb]ackup.rdl
+*- [Bb]ackup ([0-9]).rdl
+*- [Bb]ackup ([0-9][0-9]).rdl
+
+# Microsoft Fakes
+FakesAssemblies/
+
+# GhostDoc plugin setting file
+*.GhostDoc.xml
+
+# Node.js Tools for Visual Studio
+.ntvs_analysis.dat
+node_modules/
+
+# Visual Studio 6 build log
+*.plg
+
+# Visual Studio 6 workspace options file
+*.opt
+
+# Visual Studio 6 auto-generated workspace file (contains which files were open etc.)
+*.vbw
+
+# Visual Studio 6 auto-generated project file (contains which files were open etc.)
+*.vbp
+
+# Visual Studio 6 workspace and project file (working project files containing files to include in project)
+*.dsw
+*.dsp
+
+# Visual Studio 6 technical files
+*.ncb
+*.aps
+
+# Visual Studio LightSwitch build output
+**/*.HTMLClient/GeneratedArtifacts
+**/*.DesktopClient/GeneratedArtifacts
+**/*.DesktopClient/ModelManifest.xml
+**/*.Server/GeneratedArtifacts
+**/*.Server/ModelManifest.xml
+_Pvt_Extensions
+
+# Paket dependency manager
+.paket/paket.exe
+paket-files/
+
+# FAKE - F# Make
+.fake/
+
+# CodeRush personal settings
+.cr/personal
+
+# Python Tools for Visual Studio (PTVS)
+__pycache__/
+*.pyc
+
+# Cake - Uncomment if you are using it
+# tools/**
+# !tools/packages.config
+
+# Tabs Studio
+*.tss
+
+# Telerik's JustMock configuration file
+*.jmconfig
+
+# BizTalk build output
+*.btp.cs
+*.btm.cs
+*.odx.cs
+*.xsd.cs
+
+# OpenCover UI analysis results
+OpenCover/
+
+# Azure Stream Analytics local run output
+ASALocalRun/
+
+# MSBuild Binary and Structured Log
+*.binlog
+
+# NVidia Nsight GPU debugger configuration file
+*.nvuser
+
+# MFractors (Xamarin productivity tool) working folder
+.mfractor/
+
+# Local History for Visual Studio
+.localhistory/
+
+# Visual Studio History (VSHistory) files
+.vshistory/
+
+# BeatPulse healthcheck temp database
+healthchecksdb
+
+# Backup folder for Package Reference Convert tool in Visual Studio 2017
+MigrationBackup/
+
+# Ionide (cross platform F# VS Code tools) working folder
+.ionide/
+
+# Fody - auto-generated XML schema
+FodyWeavers.xsd
+
+# VS Code files for those working on multiple tools
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+*.code-workspace
+
+# Local History for Visual Studio Code
+.history/
+
+# Windows Installer files from build outputs
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# JetBrains Rider
+*.sln.iml
+
+##
+## Visual studio for Mac
+##
+
+
+# globs
+Makefile.in
+*.userprefs
+*.usertasks
+config.make
+config.status
+aclocal.m4
+install-sh
+autom4te.cache/
+*.tar.gz
+tarballs/
+test-results/
+
+# Mac bundle stuff
+*.dmg
+*.app
+
+# content below from: https://github.com/github/gitignore/blob/master/Global/macOS.gitignore
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# content below from: https://github.com/github/gitignore/blob/master/Global/Windows.gitignore
+# Windows thumbnail cache files
+Thumbs.db
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+*.db
+.idea/

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.8
+
+WORKDIR /app
+
+COPY requirements.txt requirements.txt
+RUN pip3 install -r requirements.txt
+
+COPY . .
+
+CMD ["sh", "-c", "python3 main.py"]

dbimutils.py

@@ -0,0 +1,69 @@
+# DanBooru IMage Utility functions
+
+import cv2
+import numpy as np
+from PIL import Image
+import requests
+import io
+
+def smart_imread(img, flag=cv2.IMREAD_UNCHANGED):
+    if img.endswith(".gif"):
+        img = Image.open(img)
+        img = img.convert("RGB")
+        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+    else:
+        img = cv2.imread(img, flag)
+    return img
+
+
+def smart_24bit(img):
+    if img.dtype is np.dtype(np.uint16):
+        img = (img / 257).astype(np.uint8)
+
+    if len(img.shape) == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    elif img.shape[2] == 4:
+        trans_mask = img[:, :, 3] == 0
+        img[trans_mask] = [255, 255, 255, 255]
+        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
+    return img
+
+
+def make_square(img, target_size):
+    old_size = img.shape[:2]
+    desired_size = max(old_size)
+    desired_size = max(desired_size, target_size)
+
+    delta_w = desired_size - old_size[1]
+    delta_h = desired_size - old_size[0]
+    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
+    left, right = delta_w // 2, delta_w - (delta_w // 2)
+
+    color = [255, 255, 255]
+    new_im = cv2.copyMakeBorder(
+        img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
+    )
+    return new_im
+
+
+def smart_resize(img, size):
+    # Assumes the image has already gone through make_square
+    if img.shape[0] > size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
+    elif img.shape[0] < size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
+    return img
+
+headers = {
+    "User-Agent": "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:47.0) Gecko/20100101 Firefox/47.0"
+}
+
+def read_img_from_url(image):
+    if isinstance(image, str):
+        if image.startswith("http"):
+            res = requests.get(image, headers=headers)
+            rawimage = Image.open(io.BytesIO(res.content))
+            return rawimage
+        else:
+            raise Exception("Invalid image type")
+    return image

img_label.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import functools
+import io
+import urllib
+from typing import Tuple, List, Any
+
+import huggingface_hub
+import onnxruntime as rt
+import pandas as pd
+import numpy as np
+import PIL.Image
+import requests
+
+import dbimutils
+import piexif
+import piexif.helper
+from urllib.request import urlopen
+
+import model
+
+HF_TOKEN = ""
+SWIN_MODEL_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
+CONV_MODEL_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
+CONV2_MODEL_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
+VIT_MODEL_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
+MODEL_FILENAME = "model.onnx"
+LABEL_FILENAME = "selected_tags.csv"
+
+
+def change_model(model_name):
+    global loaded_models
+
+    if model_name == "SwinV2":
+        model = load_model(SWIN_MODEL_REPO, MODEL_FILENAME)
+    elif model_name == "ConvNext":
+        model = load_model(CONV_MODEL_REPO, MODEL_FILENAME)
+    elif model_name == "ConvNextV2":
+        model = load_model(CONV2_MODEL_REPO, MODEL_FILENAME)
+    elif model_name == "ViT":
+        model = load_model(VIT_MODEL_REPO, MODEL_FILENAME)
+
+    loaded_models[model_name] = model
+    return loaded_models[model_name]
+
+
+def load_model(model_repo: str, model_filename: str) -> rt.InferenceSession:
+    path = huggingface_hub.hf_hub_download(
+        model_repo, model_filename, use_auth_token=HF_TOKEN
+    )
+    model = rt.InferenceSession(path)
+    return model
+
+
+def load_labels() -> tuple[list[Any], list[Any], list[Any], list[Any]]:
+    path = huggingface_hub.hf_hub_download(
+        CONV2_MODEL_REPO, LABEL_FILENAME, use_auth_token=HF_TOKEN
+    )
+    df = pd.read_csv(path)
+
+    tag_names = df["name"].tolist()
+    rating_indexes = list(np.where(df["category"] == 9)[0])
+    general_indexes = list(np.where(df["category"] == 0)[0])
+    character_indexes = list(np.where(df["category"] == 4)[0])
+    return tag_names, rating_indexes, general_indexes, character_indexes
+
+
+def predict(
+        image: PIL.Image.Image,
+        model_name: str,
+        general_threshold: float,
+        character_threshold: float,
+        tag_names: list[str],
+        rating_indexes: list[np.int64],
+        general_indexes: list[np.int64],
+        character_indexes: list[np.int64],
+):
+    global loaded_models
+
+    if isinstance(image, str):
+        rawimage = dbimutils.read_img_from_url(image)
+    elif isinstance(image, PIL.Image.Image):
+        rawimage = image
+    else:
+        raise Exception("Invalid image type")
+
+    image = rawimage
+
+    model = loaded_models[model_name]
+    if model is None:
+        model = change_model(model_name)
+
+    _, height, width, _ = model.get_inputs()[0].shape
+
+    # Alpha to white
+    image = image.convert("RGBA")
+    new_image = PIL.Image.new("RGBA", image.size, "WHITE")
+    new_image.paste(image, mask=image)
+    image = new_image.convert("RGB")
+    image = np.asarray(image)
+
+    # PIL RGB to OpenCV BGR
+    image = image[:, :, ::-1]
+
+    image = dbimutils.make_square(image, height)
+    image = dbimutils.smart_resize(image, height)
+    image = image.astype(np.float32)
+    image = np.expand_dims(image, 0)
+
+    input_name = model.get_inputs()[0].name
+    label_name = model.get_outputs()[0].name
+    probs = model.run([label_name], {input_name: image})[0]
+
+    labels = list(zip(tag_names, probs[0].astype(float)))
+
+    # First 4 labels are actually ratings: pick one with argmax
+    ratings_names = [labels[i] for i in rating_indexes]
+    rating = dict(ratings_names)
+
+    # Then we have general tags: pick any where prediction confidence > threshold
+    general_names = [labels[i] for i in general_indexes]
+    general_res = [x for x in general_names if x[1] > general_threshold]
+    general_res = dict(general_res)
+
+    # Everything else is characters: pick any where prediction confidence > threshold
+    character_names = [labels[i] for i in character_indexes]
+    character_res = [x for x in character_names if x[1] > character_threshold]
+    character_res = dict(character_res)
+
+    b = dict(sorted(general_res.items(), key=lambda item: item[1], reverse=True))
+    a = (
+        ", ".join(list(b.keys()))
+        .replace("_", " ")
+        .replace("(", "\(")
+        .replace(")", "\)")
+    )
+    c = ", ".join(list(b.keys()))
+
+    items = rawimage.info
+    geninfo = ""
+
+    if "exif" in rawimage.info:
+        exif = piexif.load(rawimage.info["exif"])
+        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
+        try:
+            exif_comment = piexif.helper.UserComment.load(exif_comment)
+        except ValueError:
+            exif_comment = exif_comment.decode("utf8", errors="ignore")
+
+        items["exif comment"] = exif_comment
+        geninfo = exif_comment
+
+        for field in [
+            "jfif",
+            "jfif_version",
+            "jfif_unit",
+            "jfif_density",
+            "dpi",
+            "exif",
+            "loop",
+            "background",
+            "timestamp",
+            "duration",
+        ]:
+            items.pop(field, None)
+
+    geninfo = items.get("parameters", geninfo)
+
+    for key, text in items.items():
+        print(key)
+        print(text)
+
+    print("geninfo", geninfo)
+    print("a", a)
+    print("c", c)
+    print("rating", rating)
+    print("character_res", character_res)
+    print("general_res", general_res)
+
+    character_res = list(filter(lambda x: x['confidence'] > 0.4, [{'tag': tag, 'confidence': score}
+                                                                  for tag, score in character_res.items()]))
+
+    general_res = list(filter(lambda x: x['confidence'] > 0.4, [{'tag': tag, 'confidence': score}
+                                                                for tag, score in general_res.items()]))
+
+    return {'a': a, 'c': c, 'rating': rating, 'character_res': character_res, 'general_res': general_res}
+
+
+def label_img(
+        image: PIL.Image.Image | str,
+        model: str,
+        # model: (["SwinV2", "ConvNext", "ConvNextV2", "ViT"], value="ConvNextV2", label="Model"),
+        l_score_general_threshold: float,
+        l_score_character_threshold: float,
+):
+    if isinstance(image, str) and image.startswith("http"):
+        image = dbimutils.read_img_from_url(image)
+
+    global loaded_models
+    loaded_models = {"SwinV2": None, "ConvNext": None, "ConvNextV2": None, "ViT": None}
+
+    change_model("ConvNextV2")
+
+    tag_names, rating_indexes, general_indexes, character_indexes = load_labels()
+
+    func = functools.partial(
+        predict,
+        tag_names=tag_names,
+        rating_indexes=rating_indexes,
+        general_indexes=general_indexes,
+        character_indexes=character_indexes,
+    )
+
+    return func(
+        image=image, model_name=model,
+        general_threshold=l_score_general_threshold,
+        character_threshold=l_score_character_threshold,
+    )
+
+
+def write_image_tag(img_id: int, is_valid: bool, tags: List[model.ImageTag], callback_url: str):
+    model.ImageScanCallbackRequest(img_id=img_id, is_valid=is_valid, tags=tags)
+
+
+if __name__ == "__main__":
+    score_slider_step = 0.05
+    score_general_threshold = 0.35
+    score_character_threshold = 0.85
+
+    ret = label_img(
+        image='https://pub-9747017e9ec54620bfbe2385f14fe4d7.r2.dev/cnGirlYcy_v10_people_network_nannansleep/cnGirlYcy_v10_people_network_nannansleep_r_1679670778_0.png',
+        model="SwinV2",
+        l_score_general_threshold=score_general_threshold,
+        l_score_character_threshold=score_character_threshold,
+    )
+    print(ret)

img_nsfw.py

@@ -0,0 +1,83 @@
+concepts = ['sexual', 'nude', 'sex', '18+', 'naked', 'nsfw', 'porn', 'dick', 'vagina', 'naked person (approximation)',
+            'explicit content', 'uncensored', 'fuck', 'nipples', 'nipples (approximation)', 'naked breasts', 'areola']
+special_concepts = ["small girl (approximation)", "young child", "young girl"]
+
+import dbimutils
+
+
+def init_nsfw_pipe():
+    import torch
+    from diffusers import StableDiffusionPipeline
+    from torch import nn
+
+    # make sure you're logged in with `huggingface-cli login`
+    pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16",
+                                                   torch_dtype=torch.float16)
+    pipe = pipe.to('cuda')
+
+    def cosine_distance(image_embeds, text_embeds):
+        normalized_image_embeds = nn.functional.normalize(image_embeds)
+        normalized_text_embeds = nn.functional.normalize(text_embeds)
+        return torch.mm(normalized_image_embeds, normalized_text_embeds.t())
+
+    @torch.no_grad()
+    def forward_ours(self, clip_input, images):
+        pooled_output = self.vision_model(clip_input)[1]  # pooled_output
+        image_embeds = self.visual_projection(pooled_output)
+
+        special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().numpy()
+        cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().numpy()
+
+        result = []
+        batch_size = image_embeds.shape[0]
+        for i in range(batch_size):
+            result_img = {"special_scores": {}, "special_care": [], "concept_scores": {}, "bad_concepts": []}
+
+            # increase this value to create a stronger `nfsw` filter
+            # at the cost of increasing the possibility of filtering benign images
+            adjustment = 0.0
+
+            for concet_idx in range(len(special_cos_dist[0])):
+                concept_cos = special_cos_dist[i][concet_idx]
+                concept_threshold = self.special_care_embeds_weights[concet_idx].item()
+                result_img["special_scores"][concet_idx] = round(concept_cos - concept_threshold + adjustment, 3)
+                if result_img["special_scores"][concet_idx] > 0:
+                    result_img["special_care"].append({"tag": special_concepts[concet_idx],
+                                                       "confidence": result_img["special_scores"][concet_idx]})
+                    adjustment = 0.01
+                    print("Special concept matched:", special_concepts[concet_idx])
+
+            for concet_idx in range(len(cos_dist[0])):
+                concept_cos = cos_dist[i][concet_idx]
+                concept_threshold = self.concept_embeds_weights[concet_idx].item()
+                result_img["concept_scores"][concet_idx] = round(concept_cos - concept_threshold + adjustment, 3)
+                # print("no-special", concet_idx, concepts[concet_idx], concept_threshold, round(concept_cos - concept_threshold + adjustment, 3))
+                if result_img["concept_scores"][concet_idx] > 0:
+                    result_img["bad_concepts"].append({"tag": concepts[concet_idx],
+                                                       "confidence": result_img["concept_scores"][concet_idx]})
+                    print("NSFW concept found:", concepts[concet_idx])
+
+            special_tags = list(filter(lambda x: x['confidence'] > 0.4, result_img['special_care']))
+            bad_tags = list(filter(lambda x: x['confidence'] > 0.4, result_img['bad_concepts']))
+
+            result.append({"special_tags": special_tags,
+                           "bad_tags": bad_tags, })
+
+        return images, result
+
+    from functools import partial
+    pipe.safety_checker.forward = partial(forward_ours, self=pipe.safety_checker)
+
+    return pipe
+
+
+def check_nsfw(img, pipe):
+    if isinstance(img, str):
+        img = dbimutils.read_img_from_url(img)
+        safety_checker_input = pipe.feature_extractor(images=img, return_tensors="pt").to("cuda")
+    else:
+        safety_checker_input = pipe.feature_extractor(images=img, return_tensors="pt").to("cuda")
+    from torch.cuda.amp import autocast
+    with autocast():
+        _, nsfw_tags = pipe.safety_checker.forward(clip_input=safety_checker_input.pixel_values, images=img)
+        return nsfw_tags

main.py

@@ -0,0 +1,73 @@
+from typing import Union, List
+
+import requests
+import uvicorn
+from fastapi import BackgroundTasks, FastAPI
+import img_label
+from img_nsfw import init_nsfw_pipe, check_nsfw
+import model
+
+app = FastAPI()
+
+
+def write_scan_img_result(image_id: int, scans: List[int], img: str, callback: str):
+    score_general_threshold = 0.35
+    score_character_threshold = 0.85
+
+    nsfw_tags = []
+    img_tags = []
+    if 0 in scans:
+        nsfw_tags = check_nsfw(img, pipe)
+    if 1 in scans:
+        img_tags = img_label.label_img(
+            image=img,
+            model="SwinV2",
+            l_score_general_threshold=score_general_threshold,
+            l_score_character_threshold=score_character_threshold,
+        )['general_res']
+    print(nsfw_tags)
+    print(img_tags)
+    img_tags = list(map(lambda x: model.ImageTag(tag=x['tag'], confidence=x['confidence']), img_tags))
+
+    callBackReq = model.ImageScanCallbackRequest(id=image_id, isValid=True, tags=img_tags)
+    try:
+
+        requests.post(callback, json=callBackReq.dict())
+    except Exception as ex:
+        print(ex)
+
+    nsfw_tags = map(lambda x: model.ImageScanTag(type="Moderation", confidence=x['confidence']), nsfw_tags)
+
+    ret = model.ImageScanResponse(ok=True, error="", deleted=False, blockedFor=[], tags=nsfw_tags)
+    return ret
+
+
+def write_scan_model_result(model_name: str, callback: str):
+    pass
+
+
+# @app.post("/model-scan")
+# async def send_notification(email: str, background_tasks: BackgroundTasks):
+#     background_tasks.add_task(write_scan_model_result, email, callback="")
+#     return {"message": "Notification sent in the background"}
+
+
+@app.post("/image-scan")
+async def image_scan_handler(req: model.ImageScanRequest, background_tasks: BackgroundTasks):
+    if not req.wait:
+        background_tasks.add_task(write_scan_img_result,
+                                  image_id=req.imageId,
+                                  scans=req.scans,
+                                  img=req.url, callback=req.callbackUrl)
+        return model.ImageScanResponse(ok=True, error="", deleted=False, blockedFor=[], tags=[])
+    else:
+        ret = write_scan_img_result(image_id=req.imageId,
+                                    scans=req.scans,
+                                    img=req.url, callback=req.callbackUrl)
+        return ret
+
+
+if __name__ == "__main__":
+    global pipe
+    pipe = init_nsfw_pipe()
+    uvicorn.run(app, host="0.0.0.0", port=6006)

model.py

@@ -0,0 +1,48 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class ImageScanRequest(BaseModel):
+    imageId: int
+    url: str
+    wait: bool
+    scans: List[int]
+    callbackUrl: str
+
+
+class ImageScanTag(BaseModel):
+    type: str
+    name: str
+
+
+class ImageScanResponse(BaseModel):
+    ok: bool
+    error: str
+    deleted: bool
+    blockedFor: List[str]
+    tags: List[ImageScanTag]
+
+
+class ImageTag(BaseModel):
+    tag: str
+    id: Optional[int]
+    confidence: int
+
+
+class ImageScanCallbackRequest(BaseModel):
+    id: int
+    isValid: bool
+    tags: List[ImageTag]
+
+
+class ModelScanRequest(BaseModel):
+    callbackUrl: str
+    fileUrl: str
+    lowPriority: bool
+    tasks: List[str]
+
+
+class ModelScanResponse(BaseModel):
+    ok: bool
+    error: str

requirements.txt

@@ -0,0 +1,15 @@
+uvicorn[standard]
+fastapi
+huggingface_hub
+pandas
+Pillow
+onnxruntime
+numpy
+piexif
+opencv-python-headless
+diffusers
+transformers
+scipy
+ftfy
+torch
+requests