Metadata-Version: 2.4
Name: omni-anime-tagger
Version: 1.0.0
Summary: A simple tagger to tag your medias by user-defined open-source AI model, especially working for Hydrus Network.
License-Expression: GPL-2.0-or-later
License-File: LICENSE
Requires-Python: >=3.11
Requires-Dist: deepdanbooru
Requires-Dist: dotenv>=0.9.9
Requires-Dist: hydrus-api==4.0.0
Requires-Dist: loguru>=0.7.3
Requires-Dist: nvidia-cuda-runtime-cu12==12.1.105
Requires-Dist: nvidia-cudnn-cu12==9.1.0.70
Requires-Dist: onnxruntime-gpu==1.20.1
Requires-Dist: opencv-python>=4.13.0.92
Requires-Dist: pandas>=3.0.1
Requires-Dist: platformdirs>=4.9.4
Requires-Dist: pydantic>=2.12.5
Requires-Dist: pyperclip>=1.11.0
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: questionary>=2.1.1
Requires-Dist: torch==2.5.1
Requires-Dist: typer>=0.24.1
Description-Content-Type: text/markdown

# Omni Tagger

Also Available in <a href="./README_zh_cn.md">Chinese(中文)</a>


A versatile media tagging CLI tool that tags images, animations, and videos using custom deep‑learning models (e.g., DeepDanbooru) and integrates seamlessly with Hydrus media server.

## Features

- **Multi‑format support**: Tags both images and videos by extracting frames.
- **Model‑agnostic**: Works with any ONNX‑format model; pre‑configured for DeepDanbooru.
- **Hydrus integration**: Pulls media by hash from a Hydrus instance and pushes predicted tags back.
- **Clipboard‑ready**: Reads SHA‑256 hashes directly from the clipboard for quick operations.
- **Batch processing**: Tags multiple files in one run with a progress bar.
- **Automatic namespace assignment**: Can optionally prepend category‑based namespaces (e.g., `character:`, `rating:`) to tags.
- **Rating detection**: Models that output rating scores (safe, questionable, explicit, etc.) are supported.
- **Extensible registry**: Add your own models by dropping a folder into `models/` with a `metainfo.yaml` file.

## Installation

### Prerequisites

- Python ≥ 3.11
- CUDA 12.1 (optional, for GPU acceleration)
- A running Hydrus instance (if you want to use the Hydrus integration)

### Using uv (recommended)

This project uses [uv](https://github.com/astral-sh/uv) for dependency management and packaging.

1. Clone the repository:
   ```bash
   git clone https://github.com/your‑username/omni‑tagger.git
   cd omni‑tagger
   ```

2. Install the package in editable mode:
   ```bash
   uv pip install -e .
   ```

   This will install all dependencies listed in `pyproject.toml`, including the GPU‑ready versions of PyTorch and ONNX Runtime when CUDA is available.

### Using pip

If you prefer plain pip, you can install from the local directory:

```bash
pip install -e .
```

## Usage

After installation, the CLI command `hydrus` is available.

### 1. Set up Hydrus configuration

Create a `.env` file in the project root (or in your user config directory) with the following variables:

```env
HYDRUS_API_KEY=your_access_key_here
HYDRUS_ENDPOINT=http://localhost:45869
HYDRUS_TAG_SERVICE=my tags
```

Alternatively, you can set these environment variables directly.

### 2. Register a model

Omni Tagger looks for models inside a `models/` directory (created automatically on first run). Each model must be a sub‑directory containing:

- `metainfo.yaml` – model metadata (see example below)
- the ONNX checkpoint file
- a CSV or TXT file listing the tags (one per line)

Example `metainfo.yaml`:

```yaml
name: "deepdanbooru"
checkpoint_path: "model.onnx"
tag_path: "tags.csv"
is_rating: false
num_rating: 0
height: 512
width: 512
```

Place the model folder under `models/` (e.g., `models/deepdanbooru/`). The registry is built automatically when the tool starts.

### 3. Tag media from Hydrus

The main command is `hydrus eval`. It reads SHA‑256 hashes from the clipboard (or from command‑line arguments) and tags the corresponding media.

```bash
# Copy some hashes to your clipboard, then run:
hydrus eval --threshold 0.35 --enable-auto-namespace
```

Options:
- `--threshold` – confidence threshold for accepting a tag (default: 0.35)
- `--enable-auto-namespace` – automatically prepend namespace based on tag categories
- `--rating-only` – only output rating tags, skip regular tags

You can also pass the hashes directly as arguments:

```bash
hydrus eval abc123def... xyz789...
```

### 4. List registered models

```bash
hydrus print_model
```

Shows the details of every model currently registered.

### 5. Stand‑alone tagging (without Hydrus)

You can also use the tagging engine programmatically. Import `TaggerModel` from `omni_tagger.modules.tagger` and follow the examples in the source code.

## Configuration Files

- **Hydrus config**: Located at `~/.config/omni‑tagger/hydrus.yaml` (or via environment variables). See `src/omni_tagger/config/hydrus.py` for the schema.
- **Logs**: Written to `logs/image‑renamer.log` (rotated weekly).

## Model Registry Details

The registry is built by scanning `models/` for sub‑directories that contain a `metainfo.yaml`. The YAML file must define at least:

- `checkpoint_path` – relative path to the ONNX file inside the model directory
- `tag_path` – relative path to the tag list (CSV or TXT)
- `is_rating` – boolean indicating whether the model outputs rating scores
- `num_rating` – number of rating classes (if `is_rating` is true)
- `height` / `width` – input dimensions expected by the model

## Development

### Project structure

```
src/omni_tagger/
├── __main__.py          # CLI entry point (Typer app)
├── modules/
│   ├── tagger/          # Core tagging logic (TaggerModel, ONNX wrapper)
│   └── hydrus/          # Hydrus API client
├── config/              # Configuration loaders
└── utils/               # Helper functions
```

### Running tests

Currently no formal test suite is set up. Contributions are welcome.

### Code style

The project uses `black` for formatting and `ruff` for linting. Please run them before submitting a pull request.

```bash
black .
ruff check --fix .
```

## Contributing

1. Fork the repository.
2. Create a feature branch.
3. Implement your changes with appropriate tests.
4. Ensure the code passes formatting and linting checks.
5. Submit a pull request with a clear description of the change.

## License

MIT – see the `LICENSE` file (if present) for details.

## Acknowledgments

- [DeepDanbooru](https://github.com/KichangKim/DeepDanbooru) for the pre‑trained tagging model.
- [Hydrus Network](https://hydrusnetwork.github.io/hydrus/) for the excellent media‑server API.
- [Typer](https://typer.tiangolo.com/) for the elegant CLI framework.
- [ONNX Runtime](https://onnxruntime.ai/) for efficient model inference.