File size: 3,284 Bytes
c9a5b52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | #!/usr/bin/env python3
"""Moondream 2 — пакетная генерация подписей к картинкам.
Usage:
python caption.py <image_or_dir>...
Для каждого изображения:
- проверка, не битое ли
- ресайз до 768 по большей стороне (пропорции сохраняются)
- конверт в RGB
- генерация normal caption
- сохранение в <имя_файла>.txt
"""
from transformers import AutoModelForCausalLM, AutoTokenizer
from PIL import Image, UnidentifiedImageError
import sys, os, torch
from pathlib import Path
from tqdm import tqdm
MODEL_ID = "vikhyatk/moondream2"
REVISION = "2025-06-21"
MAX_SIZE = 768
def load_image(path: str) -> Image.Image | None:
"""Загружает, конвертит в RGB, ресайзит. Возвращает None если битая."""
try:
img = Image.open(path)
img.load() # проверить, что не битая
except (UnidentifiedImageError, OSError, Exception) as e:
print(f" [SKIP] не удалось прочитать: {e}")
return None
# RGB
if img.mode != "RGB":
img = img.convert("RGB")
# Ресайз: большая сторона = MAX_SIZE
w, h = img.size
if max(w, h) > MAX_SIZE:
scale = MAX_SIZE / max(w, h)
new_w = round(w * scale)
new_h = round(h * scale)
img = img.resize((new_w, new_h), Image.LANCZOS)
return img
def process_image(model, tokenizer, img_path: str):
"""Обрабатывает одно изображение."""
path = Path(img_path)
txt_path = path.with_suffix(".txt")
if txt_path.exists():
return
img = load_image(img_path)
if img is None:
return
caption = model.caption(img, length="normal")["caption"]
txt_path.write_text(caption.strip() + "\n")
def main():
if len(sys.argv) < 2:
print(__doc__)
sys.exit(1)
# Собираем все файлы из аргументов
EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif"}
paths: list[str] = []
for arg in sys.argv[1:]:
p = Path(arg)
if p.is_dir():
paths.extend(str(f) for f in p.rglob("*") if f.suffix.lower() in EXTS)
elif p.is_file():
paths.append(arg)
else:
print(f"Не найден: {arg}")
if not paths:
print("Нет файлов для обработки.")
sys.exit(1)
if not torch.cuda.is_available():
print("Ошибка: CUDA не доступна. Нужен GPU.")
sys.exit(1)
device = "cuda:0"
print(f"Загрузка модели {MODEL_ID} (rev {REVISION}) на {device}...")
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
revision=REVISION,
trust_remote_code=True,
dtype=torch.bfloat16,
).to(device)
model = torch.compile(model, mode="reduce-overhead")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=REVISION)
print()
for p in tqdm(paths, desc="Captioning", unit="img"):
process_image(model, tokenizer, p)
print("\nГотово.")
if __name__ == "__main__":
main()
|