|
import os |
|
from pathlib import Path |
|
import fitz |
|
from PIL import Image |
|
import pytesseract |
|
from transformers import BlipProcessor, BlipForConditionalGeneration |
|
import io |
|
import torch |
|
import gradio as gr |
|
|
|
|
|
OUTPUT_DIR = Path("outputs") |
|
OUTPUT_DIR.mkdir(exist_ok=True) |
|
|
|
def pdf_to_images(pdf_path): |
|
""" |
|
Convert PDF pages to appropriately sized images |
|
""" |
|
try: |
|
|
|
pdf_document = fitz.open(pdf_path) |
|
images = [] |
|
|
|
for page_num in range(len(pdf_document)): |
|
page = pdf_document[page_num] |
|
|
|
|
|
rect = page.rect |
|
width = rect.width |
|
height = rect.height |
|
|
|
|
|
|
|
zoom = 2000 / max(width, height) |
|
|
|
|
|
mat = fitz.Matrix(zoom, zoom) |
|
|
|
|
|
pix = page.get_pixmap(matrix=mat) |
|
|
|
|
|
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) |
|
|
|
|
|
image_path = OUTPUT_DIR / f"page_{page_num + 1}.png" |
|
img.save(image_path, "PNG") |
|
images.append((image_path, img)) |
|
|
|
pdf_document.close() |
|
return images |
|
except Exception as e: |
|
print(f"Error converting PDF to images: {str(e)}") |
|
return [] |
|
|
|
def extract_text_from_image(image): |
|
""" |
|
Extract text from an image using OCR |
|
""" |
|
try: |
|
text = pytesseract.image_to_string(image) |
|
return text.strip() |
|
except Exception as e: |
|
print(f"Error during OCR: {str(e)}") |
|
return "" |
|
|
|
def analyze_image(image_path): |
|
""" |
|
Analyze image content using BLIP model for image captioning |
|
""" |
|
try: |
|
|
|
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") |
|
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base") |
|
|
|
|
|
image = Image.open(image_path).convert('RGB') |
|
inputs = processor(image, return_tensors="pt") |
|
|
|
|
|
with torch.no_grad(): |
|
outputs = model.generate(**inputs) |
|
caption = processor.decode(outputs[0], skip_special_tokens=True) |
|
|
|
return caption |
|
except Exception as e: |
|
print(f"Error during image analysis: {str(e)}") |
|
return "Image content could not be analyzed." |
|
|
|
def process_pdf(pdf_path, output_txt_path): |
|
""" |
|
Main function to process the PDF and generate output |
|
""" |
|
|
|
print("Converting PDF to images...") |
|
images = pdf_to_images(pdf_path) |
|
|
|
if not images: |
|
print("No images were generated from the PDF.") |
|
return |
|
|
|
|
|
with open(output_txt_path, 'w', encoding='utf-8') as f: |
|
f.write(f"Analysis of {os.path.basename(pdf_path)}\n") |
|
f.write("=" * 50 + "\n\n") |
|
|
|
|
|
for page_num, (image_path, image) in enumerate(images, 1): |
|
print(f"Processing page {page_num}...") |
|
|
|
|
|
f.write(f"Page {page_num}\n") |
|
f.write("-" * 30 + "\n\n") |
|
|
|
|
|
text = extract_text_from_image(image) |
|
if text: |
|
f.write("Extracted Text:\n") |
|
f.write(text) |
|
f.write("\n\n") |
|
else: |
|
f.write("No text could be extracted from this page.\n\n") |
|
|
|
|
|
description = analyze_image(image_path) |
|
f.write("Image Description:\n") |
|
f.write(f"{description}\n") |
|
f.write("\n" + "=" * 50 + "\n\n") |
|
|
|
print(f"Processing complete. Results saved to {output_txt_path}") |
|
|
|
def process_uploaded_pdf(pdf_file): |
|
if pdf_file is None: |
|
return "Please upload a PDF file." |
|
|
|
output_txt = OUTPUT_DIR / "analysis_results.txt" |
|
process_pdf(pdf_file.name, output_txt) |
|
|
|
|
|
with open(output_txt, 'r', encoding='utf-8') as f: |
|
results = f.read() |
|
|
|
return results |
|
|
|
|
|
interface = gr.Interface( |
|
fn=process_uploaded_pdf, |
|
inputs=gr.File(label="Upload PDF"), |
|
outputs=gr.Textbox(label="Analysis Results"), |
|
title="PDF Analyzer", |
|
description="Upload a PDF file to extract text and analyze images." |
|
) |
|
|
|
interface.launch() |