AI-BOOK / app.py
ginipick's picture
Update app.py
01a933c verified
raw
history blame
10.1 kB
import os
import gradio as gr
import shutil
import tempfile
import uuid
from pathlib import Path
import json
from PIL import Image
import fitz # PyMuPDF for PDF handling
# Constants
TEMP_DIR = "temp"
UPLOAD_DIR = os.path.join(TEMP_DIR, "uploads")
OUTPUT_DIR = os.path.join(TEMP_DIR, "output")
THUMBS_DIR = os.path.join(OUTPUT_DIR, "thumbs")
# Ensure directories exist
for dir_path in [TEMP_DIR, UPLOAD_DIR, OUTPUT_DIR, THUMBS_DIR]:
os.makedirs(dir_path, exist_ok=True)
def create_thumbnail(image_path, output_path, size=(300, 300)):
"""Create a thumbnail from an image."""
try:
with Image.open(image_path) as img:
img.thumbnail(size, Image.LANCZOS)
img.save(output_path)
return output_path
except Exception as e:
print(f"Error creating thumbnail: {e}")
return None
def process_pdf(pdf_path, session_id):
"""Extract pages from a PDF and save as images with thumbnails."""
pages_info = []
output_folder = os.path.join(OUTPUT_DIR, session_id)
thumbs_folder = os.path.join(THUMBS_DIR, session_id)
os.makedirs(output_folder, exist_ok=True)
os.makedirs(thumbs_folder, exist_ok=True)
try:
# Open the PDF
pdf_document = fitz.open(pdf_path)
# Process each page
for page_num, page in enumerate(pdf_document):
# Render page to an image with a higher resolution
pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
image_path = os.path.join(output_folder, f"page_{page_num + 1}.png")
pix.save(image_path)
# Create thumbnail
thumb_path = os.path.join(thumbs_folder, f"thumb_{page_num + 1}.png")
create_thumbnail(image_path, thumb_path)
# Add page info
pages_info.append({
"src": os.path.join("output", session_id, f"page_{page_num + 1}.png"),
"thumb": os.path.join("thumbs", session_id, f"thumb_{page_num + 1}.png"),
"title": f"Page {page_num + 1}"
})
return pages_info
except Exception as e:
print(f"Error processing PDF: {e}")
return []
def process_images(image_paths, session_id):
"""Process uploaded images and create thumbnails."""
pages_info = []
output_folder = os.path.join(OUTPUT_DIR, session_id)
thumbs_folder = os.path.join(THUMBS_DIR, session_id)
os.makedirs(output_folder, exist_ok=True)
os.makedirs(thumbs_folder, exist_ok=True)
for i, img_path in enumerate(image_paths):
try:
# Copy original image to output folder
dest_path = os.path.join(output_folder, f"image_{i + 1}.png")
shutil.copy(img_path, dest_path)
# Create thumbnail
thumb_path = os.path.join(thumbs_folder, f"thumb_{i + 1}.png")
create_thumbnail(img_path, thumb_path)
# Add page info
pages_info.append({
"src": os.path.join("output", session_id, f"image_{i + 1}.png"),
"thumb": os.path.join("thumbs", session_id, f"thumb_{i + 1}.png"),
"title": f"Page {i + 1}"
})
except Exception as e:
print(f"Error processing image {img_path}: {e}")
return pages_info
def create_flipbook_from_pdf(pdf_file, view_mode="webgl", skin="light"):
"""Create a flipbook from uploaded PDF."""
try:
session_id = str(uuid.uuid4())
pages_info = []
if pdf_file is not None:
# Save PDF to temp directory
pdf_path = os.path.join(UPLOAD_DIR, f"{session_id}.pdf")
with open(pdf_path, "wb") as f:
f.write(pdf_file)
# Process PDF
pages_info = process_pdf(pdf_path, session_id)
else:
return """<div style="color: red; padding: 20px;">Please upload a PDF file.</div>"""
if not pages_info:
return """<div style="color: red; padding: 20px;">Failed to process the uploaded PDF. Please try again.</div>"""
# Create and return HTML for the flipbook
flipbook_html = generate_flipbook_html(pages_info, session_id, view_mode, skin)
return flipbook_html
except Exception as e:
print(f"Error creating flipbook from PDF: {e}")
return f"""<div style="color: red; padding: 20px;">An error occurred: {str(e)}</div>"""
def create_flipbook_from_images(images, view_mode="webgl", skin="light"):
"""Create a flipbook from uploaded images."""
try:
session_id = str(uuid.uuid4())
pages_info = []
if images is not None and len(images) > 0:
# Process images
pages_info = process_images(images, session_id)
else:
return """<div style="color: red; padding: 20px;">Please upload at least one image.</div>"""
if not pages_info:
return """<div style="color: red; padding: 20px;">Failed to process the uploaded images. Please try again.</div>"""
# Create and return HTML for the flipbook
flipbook_html = generate_flipbook_html(pages_info, session_id, view_mode, skin)
return flipbook_html
except Exception as e:
print(f"Error creating flipbook from images: {e}")
return f"""<div style="color: red; padding: 20px;">An error occurred: {str(e)}</div>"""
def generate_flipbook_html(pages_info, session_id, view_mode, skin):
"""Generate HTML for the flipbook."""
# Convert pages_info to JSON for JavaScript
pages_json = json.dumps(pages_info)
# Create a custom ID for this flipbook
flipbook_id = f"flipbook_{session_id}"
# HTML template with embedded CSS and JavaScript
html = f"""
<!DOCTYPE html>
<html>
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>3D Flipbook</title>
<style>
#flipbook-container {{
width: 100%;
height: 600px;
margin: 0 auto;
position: relative;
}}
body, html {{
margin: 0;
padding: 0;
height: 100%;
overflow: hidden;
}}
</style>
<link rel="stylesheet" type="text/css" href="flipbook.css">
</head>
<body>
<div id="{flipbook_id}" class="flipbook-container"></div>
<script src="flipbook.js"></script>
<script src="flipbook.webgl.js"></script>
<script src="flipbook.swipe.js"></script>
<script src="flipbook.scroll.js"></script>
<script src="flipbook.book3.js"></script>
<script>
// Initialize flipbook when page loads
document.addEventListener('DOMContentLoaded', function() {{
const options = {{
pages: {pages_json},
viewMode: '{view_mode}',
skin: '{skin}',
responsiveView: true,
singlePageMode: false,
singlePageModeIfMobile: true,
pageFlipDuration: 1,
sound: false,
thumbnailsOnStart: true,
btnThumbs: {{ enabled: true }},
btnPrint: {{ enabled: false }},
btnDownloadPages: {{ enabled: false }},
btnDownloadPdf: {{ enabled: false }}
}};
const container = document.getElementById('{flipbook_id}');
new FlipBook(container, options);
}});
</script>
</body>
</html>
"""
return html
# Define the Gradio interface
with gr.Blocks(title="3D Flipbook Viewer") as demo:
gr.Markdown("# 3D Flipbook Viewer")
gr.Markdown("Upload a PDF file or multiple images to create an interactive 3D flipbook.")
with gr.Tabs():
with gr.TabItem("PDF Upload"):
pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"])
with gr.Accordion("Advanced Settings", open=False):
pdf_view_mode = gr.Dropdown(
choices=["webgl", "3d", "2d", "swipe"],
value="webgl",
label="View Mode"
)
pdf_skin = gr.Dropdown(
choices=["light", "dark", "gradient"],
value="light",
label="Skin"
)
pdf_create_btn = gr.Button("Create Flipbook from PDF")
pdf_output = gr.HTML(label="Flipbook Output")
# Set up PDF event handler
pdf_create_btn.click(
fn=create_flipbook_from_pdf,
inputs=[pdf_file, pdf_view_mode, pdf_skin],
outputs=pdf_output
)
with gr.TabItem("Image Upload"):
images = gr.File(label="Upload Images", file_types=["image"], file_count="multiple")
with gr.Accordion("Advanced Settings", open=False):
img_view_mode = gr.Dropdown(
choices=["webgl", "3d", "2d", "swipe"],
value="webgl",
label="View Mode"
)
img_skin = gr.Dropdown(
choices=["light", "dark", "gradient"],
value="light",
label="Skin"
)
img_create_btn = gr.Button("Create Flipbook from Images")
img_output = gr.HTML(label="Flipbook Output")
# Set up image event handler
img_create_btn.click(
fn=create_flipbook_from_images,
inputs=[images, img_view_mode, img_skin],
outputs=img_output
)
# Launch the app
if __name__ == "__main__":
demo.launch()