AI-BOOK / app.py
ginipick's picture
Create app.py
c81c490 verified
raw
history blame
8.63 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(upload_type, pdf_file=None, images=None, view_mode="webgl", skin="light"):
"""Create a flipbook from uploaded PDF or images."""
try:
session_id = str(uuid.uuid4())
pages_info = []
# Process based on upload type
if upload_type == "pdf" and 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)
elif upload_type == "images" and images is not None:
# Process images
pages_info = process_images(images, session_id)
else:
return """<div style="color: red; padding: 20px;">Please upload a PDF file or images.</div>"""
if not pages_info:
return """<div style="color: red; padding: 20px;">Failed to process the uploaded file(s). 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: {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"])
pdf_create_btn = gr.Button("Create Flipbook from PDF")
with gr.TabItem("Image Upload"):
images = gr.File(label="Upload Images", file_types=["image"], file_count="multiple")
img_create_btn = gr.Button("Create Flipbook from Images")
with gr.Accordion("Advanced Settings", open=False):
view_mode = gr.Dropdown(
choices=["webgl", "3d", "2d", "swipe"],
value="webgl",
label="View Mode"
)
skin = gr.Dropdown(
choices=["light", "dark", "gradient"],
value="light",
label="Skin"
)
output = gr.HTML(label="Flipbook Output")
# Set up event handlers
pdf_create_btn.click(
fn=create_flipbook,
inputs=[gr.Textbox(value="pdf", visible=False), pdf_file, None, view_mode, skin],
outputs=output
)
img_create_btn.click(
fn=create_flipbook,
inputs=[gr.Textbox(value="images", visible=False), None, images, view_mode, skin],
outputs=output
)
# Launch the app
if __name__ == "__main__":
demo.launch()