Spaces:
Sleeping
Sleeping
| 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) |