Views
No views yet
1from PIL import Image
2import numpy as np
3from transformers import AutoModel
4import torch
5
6model = AutoModel.from_pretrained("ragavsachdeva/magiv2", trust_remote_code=True).cuda().eval()
7
8
9def read_image(path_to_image):
10 with open(path_to_image, "rb") as file:
11 image = Image.open(file).convert("L").convert("RGB")
12 image = np.array(image)
13 return image
14
15chapter_pages = ["page1.png", "page2.png", "page3.png" ...]
16character_bank = {
17 "images": ["char1.png", "char2.png", "char3.png", "char4.png" ...],
18 "names": ["Luffy", "Sanji", "Zoro", "Ussop" ...]
19}
20
21chapter_pages = [read_image(x) for x in chapter_pages]
22character_bank["images"] = [read_image(x) for x in character_bank["images"]]
23
24with torch.no_grad():
25 per_page_results = model.do_chapter_wide_prediction(chapter_pages, character_bank, use_tqdm=True, do_ocr=True)
26
27transcript = []
28for i, (image, page_result) in enumerate(zip(chapter_pages, per_page_results)):
29 model.visualise_single_image_prediction(image, page_result, f"page_{i}.png")
30 speaker_name = {
31 text_idx: page_result["character_names"][char_idx] for text_idx, char_idx in page_result["text_character_associations"]
32 }
33 for j in range(len(page_result["ocr"])):
34 if not page_result["is_essential_text"][j]:
35 continue
36 name = speaker_name.get(j, "unsure")
37 transcript.append(f"<{name}>: {page_result['ocr'][j]}")
38with open(f"transcript.txt", "w") as fh:
39 for line in transcript:
40 fh.write(line + "\n")@misc{magiv2,
title={Tails Tell Tales: Chapter-Wide Manga Transcriptions with Character Names},
author={Ragav Sachdeva and Gyungin Shin and Andrew Zisserman},
year={2024},
eprint={2408.00298},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2408.00298},
}