Update app.py

app.py

@@ -3,33 +3,86 @@ from gradio_client import Client, handle_file
 import requests
 from PIL import Image
 import io
+import fitz  # PyMuPDF
+import tempfile
+import os
 
-def predict(image, question, seed, top_p, temperature):
-    # Inicializa o cliente do Gradio
-    client = Client("deepseek-ai/Janus-Pro-7B")
-    
-    # Prepara a imagem para envio
-    if image.startswith('http'):
-        response = requests.get(image)
-        img_path = handle_file(io.BytesIO(response.content))
-    else:
-        img_path = handle_file(image)
-    
-    # Faz a predição
-    result = client.predict(
-        image=img_path,
-        question=question,
-        seed=seed,
-        top_p=top_p,
-        temperature=temperature,
-        api_name="/multimodal_understanding"
-    )
-    
-    return result
+# Função para extrair texto e imagens de um PDF
+def extract_from_pdf(pdf_path):
+    try:
+        # Abre o PDF
+        doc = fitz.open(pdf_path)
+        extracted_text = ""
+        extracted_images = []
+
+        # Itera sobre as páginas do PDF
+        for page_num in range(len(doc)):
+            page = doc.load_page(page_num)
+            
+            # Extrai texto
+            extracted_text += page.get_text()
+            
+            # Extrai imagens
+            image_list = page.get_images(full=True)
+            for img_index, img in enumerate(image_list):
+                xref = img[0]
+                base_image = doc.extract_image(xref)
+                image_bytes = base_image["image"]
+                image = Image.open(io.BytesIO(image_bytes))
+                extracted_images.append(image)
+        
+        return extracted_text, extracted_images
+    except Exception as e:
+        return f"Erro ao processar PDF: {str(e)}", []
+
+# Função principal para fazer a predição
+def predict(file, question, seed, top_p, temperature):
+    try:
+        # Verifica se o arquivo é um PDF
+        if file.endswith(".pdf"):
+            # Extrai texto e imagens do PDF
+            extracted_text, extracted_images = extract_from_pdf(file)
+            
+            # Se houver imagens, processa a primeira imagem
+            if extracted_images:
+                image = extracted_images[0]
+                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+                    image.save(tmp_file.name, format="PNG")
+                    img_path = tmp_file.name
+            else:
+                return "Nenhuma imagem encontrada no PDF."
+            
+            # Se houver texto, adiciona ao prompt
+            if extracted_text:
+                question = f"Texto extraído do PDF:\n{extracted_text}\n\nPergunta: {question}"
+        else:
+            # Se não for PDF, trata como imagem
+            if file.startswith('http'):
+                response = requests.get(file)
+                img_path = handle_file(io.BytesIO(response.content))
+            else:
+                img_path = handle_file(file)
+        
+        # Inicializa o cliente do Gradio
+        client = Client("deepseek-ai/Janus-Pro-7B")
+        
+        # Faz a predição
+        result = client.predict(
+            image=img_path,
+            question=question,
+            seed=seed,
+            top_p=top_p,
+            temperature=temperature,
+            api_name="/multimodal_understanding"
+        )
+        
+        return result
+    except Exception as e:
+        return f"Erro durante a predição: {str(e)}"
 
 # Componentes da interface
-image_input = gr.Image(label="Upload Image", type="filepath")
-question_input = gr.Textbox(label="Question", placeholder="Ask something about the image...")
+file_input = gr.File(label="Upload PDF or Image", file_types=[".pdf", ".png", ".jpg", ".jpeg"])
+question_input = gr.Textbox(label="Question", placeholder="Ask something about the file...")
 seed_slider = gr.Slider(0, 100, value=42, label="Seed")
 top_p_slider = gr.Slider(0, 1, value=0.95, label="Top-p")
 temp_slider = gr.Slider(0, 1, value=0.1, label="Temperature")
@@ -38,7 +91,7 @@ temp_slider = gr.Slider(0, 1, value=0.1, label="Temperature")
 demo = gr.Interface(
     fn=predict,
     inputs=[
-        image_input,
+        file_input,
         question_input,
         seed_slider,
         top_p_slider,
@@ -46,9 +99,10 @@ demo = gr.Interface(
     ],
     outputs=gr.Textbox(label="Answer"),
     title="Janus-Pro-7B Multimodal Demo",
-    description="Ask questions about images using the Janus-Pro-7B model",
+    description="Ask questions about PDFs or images using the Janus-Pro-7B model",
     examples=[
-        ["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png", "What's in this image?", 42, 0.95, 0.1]
+        ["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png", "What's in this image?", 42, 0.95, 0.1],
+        ["https://example.com/sample.pdf", "Summarize the text in this PDF.", 42, 0.95, 0.1]
     ]
 )