diff --git a/picture_extractor.py b/picture_extractor.py new file mode 100644 index 0000000..eca88ff --- /dev/null +++ b/picture_extractor.py @@ -0,0 +1,407 @@ +#!/usr/bin/env python3 +""" +DocTags Picture Extractor - Extract elements from DocTags and save as separate image files. + +Usage: + python picture_extractor.py --doctags output.doctags.txt --pdf document.pdf --page 1 --output pictures +""" + +import argparse +import os +import re +import sys +from io import BytesIO +from pathlib import Path +import pdf2image +from PIL import Image + +# Regular expression to extract picture location data +PICTURE_PATTERN = r'.*?(.*?)' + +def ensure_results_folder(custom_path=None): + """Create the results folder if it doesn't exist.""" + if custom_path: + results_dir = Path(custom_path) + else: + results_dir = Path("results") + + if not results_dir.exists(): + results_dir.mkdir(parents=True) + print(f"Created directory: {results_dir}") + return results_dir + +def parse_arguments(): + """Parse command line arguments.""" + results_dir = ensure_results_folder() + + parser = argparse.ArgumentParser(description='Extract pictures from DocTags format') + parser.add_argument('--doctags', '-d', type=str, required=True, + help='Path to DocTags file') + parser.add_argument('--pdf', '-p', type=str, required=True, + help='Path to original PDF file') + parser.add_argument('--page', type=int, default=1, + help='Page number in PDF (starts at 1, default: 1)') + parser.add_argument('--output', '-o', type=str, default=str(results_dir / "pictures"), + help='Output directory for extracted pictures') + parser.add_argument('--dpi', type=int, default=300, + help='DPI for PDF rendering (higher values produce larger images)') + parser.add_argument('--max-width', type=int, default=1200, + help='Maximum width of output images in pixels') + parser.add_argument('--adjust', action='store_true', + help='Try to automatically adjust scaling') + parser.add_argument('--scale', type=float, default=1.0, + help='Scaling factor for coordinates (default: 1.0)') + parser.add_argument('--scale-x', type=float, default=None, + help='X-axis scaling factor (overrides --scale)') + parser.add_argument('--scale-y', type=float, default=None, + help='Y-axis scaling factor (overrides --scale)') + parser.add_argument('--margin', type=int, default=0, + help='Add margin around extracted pictures in pixels') + parser.add_argument('--show', '-s', action='store_true', + help='Open a file browser to the output directory when done') + return parser.parse_args() + +def load_image_from_pdf(pdf_path, page_num=1, dpi=300): + """Load a specific page from PDF as an image.""" + if not os.path.exists(pdf_path): + raise FileNotFoundError(f"PDF file not found: {pdf_path}") + + print(f"Converting PDF page {page_num} to image (DPI: {dpi})...") + try: + pdf_images = pdf2image.convert_from_path( + pdf_path, + dpi=dpi, + first_page=page_num, + last_page=page_num + ) + if not pdf_images: + raise Exception(f"Could not extract page {page_num} from PDF") + return pdf_images[0] # Return the requested page + except Exception as e: + raise Exception(f"Error converting PDF to image: {e}") + +def extract_pictures_from_doctags(doctags_path): + """Parse DocTags file and extract picture elements with their coordinates.""" + if not os.path.exists(doctags_path): + raise FileNotFoundError(f"DocTags file not found: {doctags_path}") + + with open(doctags_path, 'r', encoding='utf-8') as f: + doctags_content = f.read() + + # Find all picture elements with location information + pictures = [] + picture_matches = re.finditer(PICTURE_PATTERN, doctags_content, re.DOTALL) + + for i, match in enumerate(picture_matches): + x1, y1, x2, y2, caption = match.groups() + + # Extract caption if available (remove location tags) + clean_caption = re.sub(r'', '', caption).strip() + + pictures.append({ + 'id': i + 1, + 'x1': int(x1), + 'y1': int(y1), + 'x2': int(x2), + 'y2': int(y2), + 'caption': clean_caption + }) + + return pictures + +def normalize_coordinates(pictures, image_width, image_height, grid_size=500): + """ + Normalize coordinates from the DocTags grid (0-500) to actual image dimensions. + """ + normalized_pictures = [] + + for picture in pictures: + # Clone the picture + new_picture = picture.copy() + + # Convert from grid coordinates to actual page dimensions + new_picture['x1'] = int(picture['x1'] * image_width / grid_size) + new_picture['y1'] = int(picture['y1'] * image_height / grid_size) + new_picture['x2'] = int(picture['x2'] * image_width / grid_size) + new_picture['y2'] = int(picture['y2'] * image_height / grid_size) + + normalized_pictures.append(new_picture) + + return normalized_pictures + +def auto_adjust_coordinates(pictures, image_width, image_height): + """ + Automatically adjust coordinates based on image dimensions. + """ + if not pictures: + return pictures + + # Find the maximum coordinates + max_x = max([pic['x2'] for pic in pictures]) + max_y = max([pic['y2'] for pic in pictures]) + + # If coordinates seem to be in a normalized grid (0-500 range) + if max_x <= 500 and max_y <= 500: + print(f"Detected normalized coordinates (0-500 grid)") + return normalize_coordinates(pictures, image_width, image_height) + + # Calculate appropriate scaling factors with better heuristics + if max_x > 0: + x_scale = min(image_width / max_x, 1.0) if max_x > image_width else max(image_width / max_x, 0.5) + print(f"Auto-adjusted X scale to {x_scale:.3f} (image width: {image_width}, max picture x: {max_x})") + else: + x_scale = 1.0 + + if max_y > 0: + y_scale = min(image_height / max_y, 1.0) if max_y > image_height else max(image_height / max_y, 0.5) + print(f"Auto-adjusted Y scale to {y_scale:.3f} (image height: {image_height}, max picture y: {max_y})") + else: + y_scale = 1.0 + + # Apply more aggressive adjustment if image and pictures are very different in scale + if max_x > image_width * 5 or max_x < image_width / 5: + x_scale = image_width / max_x + print(f"Major X scale adjustment to {x_scale:.3f}") + + if max_y > image_height * 5 or max_y < image_height / 5: + y_scale = image_height / max_y + print(f"Major Y scale adjustment to {y_scale:.3f}") + + # Apply the scaling to all pictures + adjusted_pictures = [] + for pic in pictures: + adjusted_pic = pic.copy() + adjusted_pic['x1'] = int(pic['x1'] * x_scale) + adjusted_pic['y1'] = int(pic['y1'] * y_scale) + adjusted_pic['x2'] = int(pic['x2'] * x_scale) + adjusted_pic['y2'] = int(pic['y2'] * y_scale) + adjusted_pictures.append(adjusted_pic) + + print(f"Applied auto-scaling: X={x_scale}, Y={y_scale}") + return adjusted_pictures + +def extract_and_save_pictures(image, pictures, output_dir, max_width=1200, margin=0): + """Extract picture regions from the image and save them as separate files.""" + # Ensure output directory exists + output_path = ensure_results_folder(output_dir) + saved_files = [] + + # Process each picture + for picture in pictures: + try: + # Add margin to coordinates if specified + x1 = max(0, picture['x1'] - margin) + y1 = max(0, picture['y1'] - margin) + x2 = min(image.width, picture['x2'] + margin) + y2 = min(image.height, picture['y2'] + margin) + + # Check if coordinates are valid + if x1 >= x2 or y1 >= y2 or x1 < 0 or y1 < 0 or x2 > image.width or y2 > image.height: + print(f"Warning: Invalid coordinates for picture {picture['id']}: ({x1},{y1})-({x2},{y2})") + continue + + # Crop the image + cropped_img = image.crop((x1, y1, x2, y2)) + + # Resize if necessary + if cropped_img.width > max_width: + ratio = max_width / cropped_img.width + new_height = int(cropped_img.height * ratio) + cropped_img = cropped_img.resize((max_width, new_height), Image.LANCZOS) + + # Generate filename + caption = picture['caption'] + if caption: + # Create a filename-safe version of the caption (first 30 chars) + safe_caption = re.sub(r'[^\w\s-]', '', caption)[:30].strip().replace(' ', '_').lower() + filename = f"picture_{picture['id']}_{safe_caption}.png" + else: + filename = f"picture_{picture['id']}.png" + + # Save the image + output_file = output_path / filename + cropped_img.save(output_file, format="PNG") + + # Create a text file with the caption if available + if caption: + caption_file = output_path / f"{output_file.stem}.txt" + with open(caption_file, 'w', encoding='utf-8') as f: + f.write(caption) + + print(f"Saved picture {picture['id']} to {output_file}") + saved_files.append(output_file) + + except Exception as e: + print(f"Error processing picture {picture['id']}: {e}") + + return saved_files + +def create_html_index(pictures, saved_files, pdf_name, page_num, output_dir): + """Create an HTML index file of all extracted pictures.""" + output_path = Path(output_dir) + index_file = output_path / "index.html" + + html = f""" + + + + Extracted Pictures from {pdf_name} - Page {page_num} + + + +

Extracted Pictures from {pdf_name} - Page {page_num}

+

Total pictures found: {len(pictures)}

+""" + + if pictures: + html += """
+""" + + for picture, file_path in zip(pictures, saved_files): + # Get relative path for the image + rel_path = file_path.name + + html += f"""
+ Picture {picture['id']} +
+

Picture {picture['id']}

+""" + + if picture['caption']: + html += f"""
{picture['caption']}
+""" + + html += f"""
Coordinates: ({picture['x1']},{picture['y1']})-({picture['x2']},{picture['y2']})
+
+
+""" + + html += """
+""" + else: + html += """
+

No pictures found on this page

+

The DocTags file doesn't contain any picture elements for this page.

+
+""" + + html += """ + +""" + + # Save the HTML file + with open(index_file, 'w', encoding='utf-8') as f: + f.write(html) + + print(f"Created index file: {index_file}") + return index_file + +def main(): + # Parse arguments + args = parse_arguments() + + # Create output directory + output_dir = ensure_results_folder(args.output) + + try: + # Extract pictures from DocTags + print(f"Extracting pictures from {args.doctags}...") + pictures = extract_pictures_from_doctags(args.doctags) + + if not pictures: + print("No picture elements found in the DocTags file.") + return + + print(f"Found {len(pictures)} picture elements.") + + # Load the image from PDF + page_image = load_image_from_pdf(args.pdf, args.page, args.dpi) + print(f"Loaded page {args.page} image: {page_image.size[0]}x{page_image.size[1]}") + + # Process coordinates + if args.adjust: + pictures = auto_adjust_coordinates(pictures, page_image.width, page_image.height) + elif args.scale != 1.0 or args.scale_x is not None or args.scale_y is not None: + # Apply manual scaling + scale_x = args.scale_x if args.scale_x is not None else args.scale + scale_y = args.scale_y if args.scale_y is not None else args.scale + + print(f"Applying manual scaling: X={scale_x}, Y={scale_y}") + for picture in pictures: + picture['x1'] = int(picture['x1'] * scale_x) + picture['y1'] = int(picture['y1'] * scale_y) + picture['x2'] = int(picture['x2'] * scale_x) + picture['y2'] = int(picture['y2'] * scale_y) + + # Extract and save pictures + saved_files = extract_and_save_pictures( + page_image, + pictures, + output_dir, + args.max_width, + args.margin + ) + + # Create HTML index + pdf_name = Path(args.pdf).stem + index_file = create_html_index(pictures, saved_files, pdf_name, args.page, output_dir) + + # Open the output directory or index file if requested + if args.show and saved_files: + import webbrowser + if sys.platform == 'darwin': # macOS + import subprocess + subprocess.run(['open', str(output_dir)]) + elif sys.platform == 'win32': # Windows + import os + os.startfile(str(output_dir)) + else: # Linux + webbrowser.open(f"file:///{os.path.abspath(index_file)}") + + except Exception as e: + print(f"Error: {e}") + import traceback + traceback.print_exc() + +if __name__ == "__main__": + main() \ No newline at end of file