Spaces:
Sleeping
Sleeping
File size: 3,185 Bytes
fe812ad 1ad252d fe812ad 18b0b96 1ad252d 18b0b96 fe812ad f8c65fd fe812ad f8c65fd fe812ad f8c65fd fe812ad 1ad252d fe812ad f8c65fd 1ad252d |
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 |
import gradio as gr
from transformers import DonutProcessor, VisionEncoderDecoderModel
import torch
from PIL import Image
import re
import sentencepiece
# Carica il modello con parametri specifici
processor = DonutProcessor.from_pretrained(
"naver-clova-ix/donut-base",
use_fast=True
)
model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()
def estrai_titolo_autore(testo):
titolo = None
autore = None
# Cerca titolo
titolo_match = re.search(r"Title[:\-–]\s*(.*?)(?:,|$)", testo, re.IGNORECASE)
if titolo_match:
titolo = titolo_match.group(1).strip()
# Cerca autore
autore_match = re.search(r"Author[:\-–]\s*(.*?)(?:,|$)", testo, re.IGNORECASE)
if autore_match:
autore = autore_match.group(1).strip()
# Fallback: prendi le prime due righe se non troviamo match
if not titolo or not autore:
linee = testo.split("\n")
if len(linee) >= 2:
titolo = titolo or linee[0].strip()
autore = autore or linee[1].strip()
return titolo or "", autore or ""
def analizza_copertina(img):
try:
# Converti l'immagine in RGB
image = img.convert("RGB")
# Prepara il prompt
prompt = "<s_docvqa><s_question>Extract book title and author from this cover.</s_question><s_answer>"
# Prepara l'input
pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
decoder_input_ids = processor.tokenizer(prompt, add_special_tokens=False, return_tensors="pt")["input_ids"].to(device)
# Genera la risposta con parametri aggiornati
outputs = model.generate(
pixel_values,
decoder_input_ids=decoder_input_ids,
max_new_tokens=100, # Aumentiamo il numero di token generati
min_length=1, # Lunghezza minima
num_beams=4, # Aumentiamo il beam search per risultati migliori
early_stopping=True,
pad_token_id=processor.tokenizer.pad_token_id,
eos_token_id=processor.tokenizer.eos_token_id,
use_cache=True,
no_repeat_ngram_size=3 # Evita ripetizioni
)
# Decodifica il risultato
testo = processor.batch_decode(outputs, skip_special_tokens=True)[0]
# Estrai titolo e autore
titolo, autore = estrai_titolo_autore(testo)
return {
"titolo": titolo,
"autore": autore,
"testo_completo": testo
}
except Exception as e:
return {
"error": str(e),
"titolo": "",
"autore": "",
"testo_completo": ""
}
# Crea l'interfaccia Gradio
iface = gr.Interface(
fn=analizza_copertina,
inputs=gr.Image(type="pil"),
outputs=gr.JSON(),
title="Estrai Titolo e Autore da una Copertina",
description="Carica un'immagine di copertina per ottenere titolo e autore del libro."
)
# Lancia l'applicazione
iface.launch(server_name="0.0.0.0", share=True) |