Skip to content

UG-Team-Data-Science/pageextractor

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

13 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Page Extractor

Extracts a page from a photo, and warps it to a rectangular image using skimage.

Features

  • Remove margin clutter from photos of pages, usually improved document processing.
  • GroundingDINO detection model integration.
  • SAM 2.1
  • Customizable text prompt, e.g. to "receipt." or "invoice."

Getting Started

Prerequisites

  • Python 3.10 or higher

Installation

Installing PyTorch Dependencies

Before installing pageextractor, install PyTorch:

pip install torch==2.4.1 torchvision==0.19.1 --extra-index-url https://download.pytorch.org/whl/cu124

Installation options

pip install -U git+https://github.com/UG-Team-Data-Science/pageextractor.git

Or:

git clone https://github.com/UG-Team-Data-Science/pageextractor && cd pageextractor
pip install -e .

Usage

from PIL import Image
from matplotlib import pyplot as plt

from pageextractor import PageExtractor

img = Image.open('example.png')
model = PageExtractor(sam_type='sam2.1_hiera_tiny', device='cuda')
mask, polygon, cropped = model.extract_page(img)

_, (ax0, ax1, ax2) = plt.subplots(1, 3, figsize=(30, 15))
ax0.imshow(img)

ax1.imshow(img)
ax1.plot(*polygon[[0,1,2,3,0]].T, 'r:')
ax1.imshow(1-mask, cmap='Blues', alpha=0.8 - 0.8*mask)

ax2.imshow(cropped)

Examples

Photo of a book page, unfortunately centerfold and page block are included

Photo of a book page, unfortunately centerfold and page block are included

Photo of a letter

Photo of a letter

Acknowledgments

This project is based on/used the following repositories:

About

Extract a page from a photo using SAM2

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages