tagger

package
v1.11.1 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Jun 24, 2026 License: AGPL-3.0 Imports: 20 Imported by: 0

Documentation

Index

Constants

View Source
const (
	DefaultModelFile    = "model.onnx"
	DefaultTagsFile     = "tags.csv"
	DefaultTextTagsFile = "tags.txt"
	// DefaultConfidenceThreshold is applied to taggers discovered on disk
	// that do not yet have a TOML entry with an explicit threshold.
	DefaultConfidenceThreshold = 0.4
)

Default filenames for a tagger subfolder. Each can be overridden in the TOML entry. When neither default is present, a lone .onnx / .csv / .txt in the folder is auto-picked.

View Source
const DefaultTopKFallback = 10

DefaultTopKFallback is applied to any category not listed in DefaultPerCategoryTopK. Custom operator-created categories land here.

Variables

View Source
var DefaultPerCategoryTopK = map[string]int{
	"character": 8,
	"copyright": 4,
	"artist":    4,
	"general":   25,
	"rating":    1,
}

DefaultPerCategoryTopK is the fallback per-category emission cap used when a tagger has no PerCategoryTopK override (and the catalog entry has no DefaultTopK either). The numbers are tuned for the WD14 / JoyTag / Camie label distributions: rare-attribution categories sit low so a noisy run can't pile dozens of characters or artists onto one row; general is generous because it carries most of the descriptive surface. The "rating" cap is 1 so the per-image rating store ends up with exactly one tag - taggers that emit several rating signals on borderline content settle deterministically on the top-scoring one instead of needing the store step to pick.

Functions

func CheckCUDAAvailable

func CheckCUDAAvailable() error

CheckCUDAAvailable always errors here; inference is compiled out.

func DispatchTargetCategories added in v1.9.2

func DispatchTargetCategories(modelPath, taggerName string) []string

DispatchTargetCategories returns the distinct destination category names the tagger's dispatch table (embedded default + on-disk overlay) routes labels into. The Configure dialog unions these with the profile's natively emitted categories so an operator can tune or disable a category the model only reaches through dispatch - e.g. wd-swinv2 routes a slice of its general labels into medium / meta / year. Empty-category entries (drops) are skipped; order is deterministic (embedded rules first, then overlay additions).

func IsAvailable

func IsAvailable(_ *config.Config) bool

func ReleaseAll added in v1.4.3

func ReleaseAll()

ReleaseAll is a no-op stub on the non-tagger build.

func ReleaseIdle added in v1.4.3

func ReleaseIdle(_ time.Duration) bool

ReleaseIdle is a no-op stub on the non-tagger build; nothing is cached when inference is compiled out.

func ResolveMinHits added in v1.7.0

func ResolveMinHits(fraction float64, frameCount int) int

ResolveMinHits returns the minimum number of pages/frames a label must score above the pre-floor on for it to survive the merge. Given the global min_hit_fraction and the frame count for this row:

  • frameCount <= 1 (static image): always 1, so the aggregation behaves identically to the legacy max-only path.
  • fraction <= 0: 1, the operator opt-out.
  • otherwise: clamp(ceil(fraction * frameCount), 2, 10).

The lower bound of 2 is what makes the gate useful - a single flicker of a noisy label across a 500-page archive shouldn't be enough to imprint it on the row. The upper bound of 10 is a guardrail so a sparse-but-real label on a long manga still has a path through.

func ResolveTopK added in v1.7.0

func ResolveTopK(overrides map[string]int, cat string) int

ResolveTopK returns the active cap for a category given a tagger's PerCategoryTopK overrides. An explicit zero in the override map disables the cap (returns 0 = uncapped). A missing key falls through to DefaultPerCategoryTopK and then DefaultTopKFallback.

func RunWithTaggers

func RunWithTaggers(_ context.Context, _ *db.DB, _ *config.Config, _ []int64, _ []TaggerStatus, _ *jobs.Manager, _ bool, _ string) (int, error)

RunWithTaggers is the no-op stub matching the tagger build signature.

func SeedTaggerInstance added in v1.5.1

func SeedTaggerInstance(name string, enabled bool, catalog *CatalogEntry) config.TaggerInstance

SeedTaggerInstance returns a fresh TaggerInstance with the given name and enabled flag, prefilled from the catalog's default_threshold, default_thresholds, and default_top_k when a non-nil entry is supplied. Used both by the discover-from-disk path and the per-row Enable handler so the same catalog-supplied defaults land in either flow without restating them. Categories absent from the catalog's DefaultTopK still resolve through DefaultPerCategoryTopK at merge time, so leaving the map nil here picks up the built-ins.

func SetBackend added in v1.7.2

func SetBackend(b Backend)

SetBackend installs the inference implementation. inproc.go's init picks the IPC backend by default and falls back to the in-proc one when MONBOORU_TAGGER_BACKEND=inproc or the IPC constructor fails.

func UnavailableReason added in v1.2.1

func UnavailableReason(_ *config.Config) string

func ValidateTaggerName added in v1.6.0

func ValidateTaggerName(name string) error

ValidateTaggerName rejects an empty or non-allowlist name. Mirrors config.ValidateGalleryName so the two operator-supplied identifiers share the same vocabulary.

func WorkerPID added in v1.7.2

func WorkerPID() (int, bool)

WorkerPID returns the live tagger-worker child's PID and true when the IPC backend currently has a worker running, or (0, false) otherwise (in-process backend, child not yet spawned, child exited). The stats panel uses this to read /proc/<pid>/smaps for the child's resident-set breakdown.

Types

type AggregateOpts added in v1.7.0

type AggregateOpts struct {
	MinHits            int
	GlobalThreshold    float32
	CategoryThresholds map[string]float64
	PerCategoryTopK    map[string]int
	DisabledCategories []string
}

AggregateOpts collects the per-merge knobs the pure aggregator needs. GlobalThreshold and CategoryThresholds gate which labels survive; PerCategoryTopK caps how many of the survivors land per category; DisabledCategories drops whole categories before any scoring. MinHits is the resolved frame-count gate from ResolveMinHits.

type AggregatedCandidate added in v1.7.0

type AggregatedCandidate struct {
	Name  string
	CatID int64
	Score float32
}

AggregatedCandidate is one (name, catID, mean-score) result that survives both the frequency gate and the per-category top-K cap. Ordering of the returned slice is undefined; callers that care (multi-tagger merge) re-key by tagKey.

func AggregateInferenceScores added in v1.7.0

func AggregateInferenceScores(perFrame [][]float32, labels []CandidateLabel, opts AggregateOpts) []AggregatedCandidate

AggregateInferenceScores applies the §2 merge to per-frame label scores. perFrame[fIdx][labelIdx] is the label's score on that frame (zero / missing = no hit). labels[labelIdx] gives the post-routing metadata. The function:

  1. drops near-zero scores (pre-floor 0.001).
  2. tracks sum + hit count per label across every frame.
  3. computes mean = sum / hits and gates on the per-category threshold and the MinHits floor.
  4. groups survivors by category, sorts by mean desc (name asc as tiebreaker) and keeps at most ResolveTopK(opts.PerCategoryTopK, CatName) per category.

Returns the survivors in no particular order.

type Backend added in v1.7.2

type Backend interface {
	// Run executes inference for one or more pre-prepared images.
	// Per-image inference errors land on the matching Result.Err so
	// the orchestrator can skip individual images without aborting
	// the whole batch.
	Run(ctx context.Context, req RunRequest) (RunResponse, error)
	// Status snapshots the warm-cache state for the operator UI.
	Status() CacheStatus
	// ReleaseIdle tears the cache down when it has been idle for at
	// least after and no run is in flight. Returns true on teardown
	// so the caller can log it.
	ReleaseIdle(after time.Duration) bool
	// ReleaseAll unconditionally tears the cache down. Called on
	// use_cuda flips and on server shutdown.
	ReleaseAll()
}

Backend is the boundary the inference loop crosses. The default implementation is the subprocess client (ipc.go); the in-process fallback lives in inproc.go for the worker subcommand and the rollback env var.

type BackendImageRequest added in v1.7.2

type BackendImageRequest struct {
	ID            int64
	FramePaths    []string
	MangaProgress bool
}

BackendImageRequest is one image's contribution to the batch: the gallery id (used for log lines + as the result key) and the already-extracted frame paths to feed inference. MangaProgress is true only on cbz rows whose page count makes the wait worth narrating; the backend then fires OnProgress per page.

type BackendImageResult added in v1.7.2

type BackendImageResult struct {
	ID   int64
	Tags map[TagKey]Scored
	Err  string
}

BackendImageResult is one image's outcome. Err carries per-image inference failures so the orchestrator can increment skipped without aborting the whole batch. Err is a string (not an error interface) because the result travels over the IPC channel via gob, and gob can't encode an interface field without Register'd concrete types.

type CacheStatus added in v1.7.2

type CacheStatus struct {
	Loaded   bool
	UseCUDA  bool
	InUse    bool
	Sessions []string
	LastUsed time.Time
}

CacheStatus is a snapshot of the tagger cache for surfacing in the operator UI. Sessions lists the loaded model names sorted so a refresh doesn't reorder the list. Returns Loaded=false when no model set is currently warm.

func Status added in v1.7.2

func Status() CacheStatus

Status reports "not loaded" since the non-tagger build never caches.

type CandidateLabel added in v1.7.0

type CandidateLabel struct {
	Name        string
	CatID       int64
	CatName     string
	Placeholder bool
}

CandidateLabel is the per-label metadata the aggregator needs after the caller has already resolved category routing (dispatch rules, single_general lifts, rating skips). One entry per label index from the model's tag vocabulary; Placeholder=true rows are dropped on sight.

type CatalogEntry added in v1.4.0

type CatalogEntry struct {
	Name              string             `json:"name"`
	Description       string             `json:"description"`
	Files             []CatalogFile      `json:"files"`
	DefaultThreshold  float64            `json:"default_threshold,omitempty"`
	DefaultThresholds map[string]float64 `json:"default_thresholds,omitempty"`
	DefaultTopK       map[string]int     `json:"default_top_k,omitempty"`
}

CatalogEntry describes one downloadable tagger: a target subfolder name under paths.model_path plus the URLs the user fetches the model and tags file from. Monbooru itself never reaches out to these URLs - the Settings → Auto-Tagger dialog only renders copy-paste curl commands so the "no automatic outbound HTTP" promise stays intact.

DefaultThreshold, DefaultThresholds, and DefaultTopK prefill TaggerInstance when an operator first enables a catalog row from Settings; all are optional. A non-zero DefaultThreshold replaces the package-wide DefaultConfidenceThreshold for that tagger; DefaultThresholds maps category name → per-category override and copies into the TaggerInstance's CategoryThresholds map; DefaultTopK maps category name → per-category cap and copies into the TaggerInstance's PerCategoryTopK map. Categories absent from DefaultTopK fall back to the built-in DefaultPerCategoryTopK table.

func LoadCatalog added in v1.4.0

func LoadCatalog(modelPath string) []CatalogEntry

LoadCatalog returns the merged tagger catalog. The embedded default catalog (two suggested taggers - WD14 SwinV2 and JoyTag) is the base; an optional <modelPath>/models.json override is applied on top so users can add or replace entries without rebuilding. Same-name entries in the override replace the default; new names append.

func (CatalogEntry) DockerCommand added in v1.4.0

func (c CatalogEntry) DockerCommand(containerName string) string

DockerCommand renders a `docker exec <container> sh -c '...'` chain that drops model files into the container's /models mount. Container name defaults to "monbooru" (matching the shipped docker-compose.yml).

func (CatalogEntry) HostCommand added in v1.4.0

func (c CatalogEntry) HostCommand() string

HostCommand renders the `mkdir + curl` chain a user runs on the host (no docker). Paths are relative to the model path.

type CatalogFile added in v1.4.0

type CatalogFile struct {
	URL      string `json:"url"`
	Filename string `json:"filename"`
}

CatalogFile is one URL-to-filename pair; Filename is the basename the file gets dropped under inside <modelPath>/<entry name>/.

type Profile added in v1.5.1

type Profile struct {
	Name string `json:"-"`

	// InputSize is the model's expected square input edge in pixels.
	// Zero means "ask ONNX at session-build time" (cache.ensure reads
	// inputs[0].Dimensions and picks the spatial axis).
	InputSize int `json:"input_size,omitempty"`

	// Layout is "nhwc" or "nchw". WD14 needs NHWC; joytag and Camie need
	// NCHW.
	Layout string `json:"layout,omitempty"`

	// Channels is "rgb" or "bgr". WD14 wants BGR (legacy TF training
	// data); the rest of the world ships RGB.
	Channels string `json:"channels,omitempty"`

	// Normalize selects the per-channel post-decode transform:
	//   "none"     - keep 0..255 byte values as float32
	//   "div255"   - rescale to 0..1 (rarely seen alone)
	//   "imagenet" - (x/255 - mean) / std with ImageNet statistics
	//   "clip"     - same shape as imagenet but with CLIP statistics
	Normalize string `json:"normalize,omitempty"`

	// Pad selects the resize/pad strategy:
	//   "white_square"      - pad to square with #FFFFFF then resize
	//   "mean_color_aspect" - resize preserving aspect into a square
	//                         canvas filled with FillColor (default
	//                         (124,116,104) - Camie's documented value)
	Pad string `json:"pad,omitempty"`

	// FillColor is the RGB triplet the "mean_color_aspect" pad fills the
	// background with. Ignored for other pad modes; defaults to
	// (124,116,104) when empty in the resolved profile.
	FillColor [3]uint8 `json:"fill_color,omitempty"`

	// Activation is "sigmoid_in_model" (WD14: outputs are already
	// probabilities in 0..1) or "logits" (joytag/Camie: monbooru applies
	// sigmoid before thresholding).
	Activation string `json:"activation,omitempty"`

	// LabelFormat selects the label-file parser:
	//   "wd14_csv"   - WD14 selected_tags.csv (id,name,category_id)
	//   "joytag_txt" - one label per line, all in `general`
	//   "camie_json" - Camie metadata JSON (idx_to_tag + tag_to_category)
	LabelFormat string `json:"label_format,omitempty"`

	// CategoryScheme drives processOne's per-label category resolution:
	//   "wd14_numeric"   - look label.categoryID up in wd14Category
	//   "single_general" - everything lands in general (joytag)
	//   "name_string"    - look label.categoryName up in tag_categories
	CategoryScheme string `json:"category_scheme,omitempty"`

	// OutputIndex selects which output tensor to read (0-based). Camie
	// ships a coarse + a refined head and only the second one is the
	// production output. Defaults to 0 for every other tagger.
	OutputIndex int `json:"output_index,omitempty"`
}

Profile captures every axis where ONNX taggers actually differ: preprocessing (input size, layout, channel order, normalisation, padding), model output shape (which output slot to read, sigmoid in-model vs raw logits), label-file format, and how labels declare their category.

A profile is resolved once per (tagger, file set) and pinned by SHA in the session cache so a sidecar edit forces a session rebuild. Empty strings mean "fall through to the previous resolution layer"; an empty final profile is invalid and surfaces at session-build time.

func ResolveProfile added in v1.5.1

func ResolveProfile(modelPath, taggerName, tagsFile string) (Profile, error)

ResolveProfile picks the runtime profile for one tagger.

Resolution order, each layer overlaying the previous (empty string = no override at this key):

  1. embedded profile_default/<taggerName>.json
  2. <modelPath>/<taggerName>/tagger.json sidecar
  3. heuristic from tagsFile extension (.csv → wd14, .txt → joytag)

A returned error means no layer could produce a valid profile; the caller surfaces that as a session-open failure rather than running with half-resolved settings.

func (Profile) EmittedCategories added in v1.5.1

func (p Profile) EmittedCategories() []string

EmittedCategories names the categories the profile is expected to produce. Used by Settings → Auto-Tagger → Configure to populate the per-category threshold dialog. The set is informational, not enforcing - a dispatch rule can route any source label into any category and the threshold lookup honours that at run time. nil from an unrecognised category scheme means "no opinion"; the caller can fall back to the canonical built-in list.

type RunRequest added in v1.7.2

type RunRequest struct {
	Cfg            *config.Config
	Taggers        []TaggerStatus
	UseCUDA        bool
	CatIDs         map[string]int64
	GeneralCatID   int64
	InferredCats   map[string]int64
	MinHitFraction float64
	Parallel       int
	Images         []BackendImageRequest
	// OnProgress fires from inside the inference loop, e.g. cbz
	// per-page status. workerIdx scopes the worker pool. Optional;
	// nil disables emission.
	OnProgress func(workerIdx int, msg string)
}

RunRequest is the per-batch payload. Each BackendImageRequest carries the pre-extracted frame paths so the backend never reads images from the gallery DB or unpacks videos / cbz archives - that work stays in the orchestrator and ports unchanged to a future subprocess split.

type RunResponse added in v1.7.2

type RunResponse struct {
	Results []BackendImageResult
}

RunResponse mirrors RunRequest.Images in order. Tags is the merged per-image score set, keyed by TagKey so the orchestrator can hand it straight to storeResults.

type Scored added in v1.7.2

type Scored struct {
	Score      float32
	TaggerName string
}

Scored carries the highest confidence seen across taggers for one TagKey plus the tagger that produced that score, so attribution survives multi-tagger merges.

type TagKey added in v1.7.2

type TagKey struct {
	Name  string
	CatID int64
}

TagKey identifies one (name, category_id) pair so multi-tagger merges never insert the same tag twice on the same image. Both the orchestrator and the worker backend use this as the merge key.

type TaggerStatus

type TaggerStatus struct {
	config.TaggerInstance
	Available bool
	Reason    string
}

TaggerStatus pairs a configured tagger with its runtime availability so the settings UI can show why each row is active or inactive.

func AvailableTaggers

func AvailableTaggers(cfg *config.Config) []TaggerStatus

AvailableTaggers lists every configured tagger as unavailable because inference is disabled at build time.

func DiscoverTaggers

func DiscoverTaggers(cfg *config.Config) []TaggerStatus

DiscoverTaggers merges tagger subfolders under paths.model_path with the configured list. The result has an entry for every on-disk folder AND every configured tagger (so leftover config is still visible after the folder vanishes). Sorted by Name.

func EnabledTaggers

func EnabledTaggers(cfg *config.Config) []TaggerStatus

EnabledTaggers returns taggers that are both enabled in config and available on disk. Returns nil on a noop build so the UI hides affordances that depend on inference. Used by surfaces that don't know which gallery a tag job is about to run on (e.g. the Settings page itself); per-gallery callers should use EnabledTaggersForGallery instead.

func EnabledTaggersForGallery added in v1.5.1

func EnabledTaggersForGallery(cfg *config.Config, gallery string) []TaggerStatus

EnabledTaggersForGallery filters EnabledTaggers down to the rows whose per-tagger Galleries list either is empty (applies to every gallery, the legacy behaviour) or contains the named gallery. Used by every per-job entry point so a tagger configured for `default` doesn't fire on `stock` and vice versa.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL