Metadata-Version: 2.4
Name: wd-tag-heatmap
Version: 0.1.0
Summary: Generate attention heatmaps for WD Tagger models - visualize which parts of an image contribute to each tag
Author: WD Tag Heatmap Contributors
Maintainer-email: Your Name <your.email@example.com>
License: MIT
Project-URL: Homepage, https://github.com/yourusername/wd-tag-heatmap
Project-URL: Documentation, https://github.com/yourusername/wd-tag-heatmap#readme
Project-URL: Repository, https://github.com/yourusername/wd-tag-heatmap
Project-URL: Bug Tracker, https://github.com/yourusername/wd-tag-heatmap/issues
Keywords: image-tagging,heatmap,attention-visualization,wd-tagger,anime,booru
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Multimedia :: Graphics
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: torchvision>=0.15.0
Requires-Dist: timm>=0.9.0
Requires-Dist: pillow>=9.0.0
Requires-Dist: numpy>=1.21.0
Requires-Dist: pandas>=1.3.0
Requires-Dist: huggingface-hub>=0.16.0
Requires-Dist: safetensors>=0.3.0
Requires-Dist: colorcet>=3.0.0
Requires-Dist: matplotlib>=3.5.0
Requires-Dist: tqdm>=4.62.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: black>=22.0.0; extra == "dev"
Requires-Dist: ruff>=0.0.260; extra == "dev"
Requires-Dist: mypy>=0.990; extra == "dev"
Dynamic: license-file

# WD Tag Heatmap

Generate attention heatmaps for WD Tagger models - visualize which parts of an image contribute to each tag prediction.

![Example heatmap grid](https://example.com/heatmap-grid.png)

## Features

- 🎯 Generate individual attention heatmaps for each detected tag
- 🖼️ Create grid visualizations of all tag heatmaps
- 📊 Export detailed tag information with confidence scores
- 🚀 Support for multiple WD Tagger v3 models
- 💻 Easy to use Python API and CLI
- 🔧 Batch processing support

## Installation

```bash
pip install wd-tag-heatmap
```

For development:
```bash
git clone https://github.com/yourusername/wd-tag-heatmap
cd wd-tag-heatmap
pip install -e ".[dev]"
```

## Quick Start

### Python API

```python
from wd_tag_heatmap import generate_tag_heatmaps

# Generate heatmaps for a single image
heatmaps, grid, labels = generate_tag_heatmaps(
    "anime_girl.jpg",
    output_dir="results",
    threshold=0.35
)

print(f"Found {len(heatmaps)} tags")
print(f"Caption: {labels.caption}")
```

### Command Line

```bash
# Process single image
wd-tag-heatmap image.jpg

# With custom settings
wd-tag-heatmap image.jpg --output results --threshold 0.5 --model SmilingWolf/wd-vit-large-tagger-v3

# Batch process
wd-tag-heatmap *.jpg --output batch_results
```

## Usage Examples

### Basic Usage

```python
from wd_tag_heatmap import generate_tag_heatmaps

# Simple usage with defaults
heatmaps, grid, labels = generate_tag_heatmaps("my_image.png")
```

### Custom Configuration

```python
from wd_tag_heatmap import generate_tag_heatmaps, AVAILABLE_MODELS

# Use a different model and threshold
heatmaps, grid, labels = generate_tag_heatmaps(
    "my_image.jpg",
    output_dir="custom_output",
    model_repo="SmilingWolf/wd-convnext-tagger-v3",
    threshold=0.5,
    save_individual_tags=True,
    save_grid=True,
    save_tags_info=True
)

# Access results programmatically
for heatmap in heatmaps[:5]:
    print(f"{heatmap.label}: {heatmap.score:.3f}")
```

### Batch Processing

```python
from wd_tag_heatmap import batch_process_images

# Process multiple images
image_paths = ["img1.jpg", "img2.png", "img3.jpg"]
batch_process_images(
    image_paths,
    output_dir="batch_results",
    threshold=0.35
)
```

### Working with Results

```python
# Get results without saving files
heatmaps, grid, labels = generate_tag_heatmaps(
    "image.jpg",
    save_individual_tags=False,
    save_grid=False,
    save_tags_info=False
)

# Access tag information
print(f"Rating tags: {labels.rating}")
print(f"Character tags: {labels.character}")
print(f"General tags: {list(labels.general.keys())[:10]}")

# Work with heatmap images
for heatmap in heatmaps[:3]:
    # heatmap.image is a PIL Image
    heatmap.image.show()
```

## Available Models

- `SmilingWolf/wd-convnext-tagger-v3`
- `SmilingWolf/wd-swinv2-tagger-v3`
- `SmilingWolf/wd-vit-tagger-v3` (default)
- `SmilingWolf/wd-vit-large-tagger-v3`
- `SmilingWolf/wd-eva02-large-tagger-v3`

## Output Structure

```
output_dir/
├── image_name_heatmap_grid.png      # Grid of all heatmaps
├── image_name_tags_info.txt         # Detailed tag information
└── individual_tags/                 # Individual heatmap images
    ├── image_name_000_1girl_0.996.png
    ├── image_name_001_solo_0.975.png
    └── ...
```

## Requirements

- Python 3.8+
- PyTorch 2.0+
- CUDA GPU recommended for faster processing

## API Reference

### `generate_tag_heatmaps()`

Main function to generate heatmaps for a single image.

**Parameters:**
- `image_path` (str): Path to input image
- `output_dir` (str): Directory for outputs (default: "heatmap_outputs")
- `model_repo` (str): Model repository (default: "SmilingWolf/wd-vit-tagger-v3")
- `threshold` (float): Confidence threshold (default: 0.35)
- `save_individual_tags` (bool): Save individual heatmaps (default: True)
- `save_grid` (bool): Save grid image (default: True)
- `save_tags_info` (bool): Save tag information (default: True)

**Returns:**
- `heatmaps` (List[Heatmap]): List of heatmap objects
- `grid` (Image): Grid visualization
- `labels` (ImageLabels): Tag information

### `batch_process_images()`

Process multiple images in batch.

**Parameters:**
- `image_paths` (List[str]): List of image paths
- Other parameters same as `generate_tag_heatmaps()`

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

## License

This project is licensed under the MIT License - see the LICENSE file for details.

## Acknowledgments

- [WD Tagger models](https://huggingface.co/SmilingWolf) by SmilingWolf
- Built with [timm](https://github.com/rwightman/pytorch-image-models)

## Citation

If you use this tool in your research, please cite:

```bibtex
@software{wd-tag-heatmap,
  title = {WD Tag Heatmap: Attention Visualization for Anime Image Tagging},
  author = {Your Name},
  year = {2024},
  url = {https://github.com/yourusername/wd-tag-heatmap}
}
```
