Engines API Reference

Complete API documentation for Doctra engines.

DocResEngine

Image restoration engine for document enhancement.

doctra.engines.image_restoration.DocResEngine

DocRes Image Restoration Engine

A wrapper around DocRes inference functionality for easy integration with Doctra's document processing pipeline.

Source code in doctra/engines/image_restoration/docres_engine.py

class DocResEngine:
    """
    DocRes Image Restoration Engine

    A wrapper around DocRes inference functionality for easy integration
    with Doctra's document processing pipeline.
    """

    SUPPORTED_TASKS = [
        'dewarping', 'deshadowing', 'appearance', 
        'deblurring', 'binarization', 'end2end'
    ]

    def __init__(
        self, 
        device: Optional[str] = None,
        use_half_precision: bool = True,
        model_path: Optional[str] = None,
        mbd_path: Optional[str] = None
    ):
        """
        Initialize DocRes Engine

        Args:
            device: Device to run on ('cuda', 'cpu', or None for auto-detect)
            use_half_precision: Whether to use half precision for inference
            model_path: Path to DocRes model checkpoint (optional, defaults to Hugging Face Hub)
            mbd_path: Path to MBD model checkpoint (optional, defaults to Hugging Face Hub)
        """
        if not DOCRES_AVAILABLE:
            raise ImportError(
                "DocRes is not available. Please install the missing dependencies:\n"
                "pip install scikit-image>=0.19.3\n\n"
                "The DocRes module is already included in this library, but requires "
                "scikit-image for image processing operations."
            )

        # Set device
        if device is None:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            requested_device = torch.device(device)
            # Check if the requested device is available
            if requested_device.type == 'cuda' and not torch.cuda.is_available():
                print(f"Warning: CUDA requested but not available. Falling back to CPU.")
                self.device = torch.device('cpu')
            else:
                self.device = requested_device

        self.use_half_precision = use_half_precision

        # Get model paths (always from Hugging Face Hub)
        try:
            self.mbd_path, self.model_path = get_model_paths(
                use_huggingface=True,
                model_path=model_path,
                mbd_path=mbd_path
            )
        except Exception as e:
            raise RuntimeError(f"Failed to get model paths: {e}")

        # Verify model files exist
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(
                f"DocRes model not found at {self.model_path}. "
                f"This may indicate a Hugging Face download failure. "
                f"Please check your internet connection and try again."
            )

        if not os.path.exists(self.mbd_path):
            raise FileNotFoundError(
                f"MBD model not found at {self.mbd_path}. "
                f"This may indicate a Hugging Face download failure. "
                f"Please check your internet connection and try again."
            )

        # Initialize model
        self._model = None
        self._initialize_model()

    def _initialize_model(self):
        """Initialize the DocRes model"""
        try:
            # Create model architecture
            self._model = restormer_arch.Restormer( 
                inp_channels=6, 
                out_channels=3, 
                dim=48,
                num_blocks=[2,3,3,4], 
                num_refinement_blocks=4,
                heads=[1,2,4,8],
                ffn_expansion_factor=2.66,
                bias=False,
                LayerNorm_type='WithBias',
                dual_pixel_task=True        
            )

            # Load model weights - always load to CPU first, then move to target device
            state = convert_state_dict(torch.load(self.model_path, map_location='cpu')['model_state'])

            self._model.load_state_dict(state)
            self._model.eval()
            self._model = self._model.to(self.device)

        except Exception as e:
            raise RuntimeError(f"Failed to initialize DocRes model: {e}")

    def restore_image(
        self, 
        image: Union[str, np.ndarray], 
        task: str = "appearance",
        save_prompts: bool = False
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        Restore a single image using DocRes

        Args:
            image: Path to image file or numpy array
            task: Restoration task to perform
            save_prompts: Whether to save intermediate prompts

        Returns:
            Tuple of (restored_image, metadata)
        """
        if task not in self.SUPPORTED_TASKS:
            raise ValueError(f"Unsupported task: {task}. Supported tasks: {self.SUPPORTED_TASKS}")

        # Load image if path provided
        if isinstance(image, str):
            if not os.path.exists(image):
                raise FileNotFoundError(f"Image not found: {image}")
            img_array = cv2.imread(image)
            if img_array is None:
                raise ValueError(f"Could not load image: {image}")
        else:
            img_array = image.copy()

        original_shape = img_array.shape

        try:
            # Handle end2end pipeline
            if task == "end2end":
                return self._run_end2end_pipeline(img_array, save_prompts)

            # Run single task
            restored_img, metadata = self._run_single_task(img_array, task, save_prompts)

            metadata.update({
                'original_shape': original_shape,
                'restored_shape': restored_img.shape,
                'task': task,
                'device': str(self.device)
            })

            return restored_img, metadata

        except Exception as e:
            raise RuntimeError(f"Image restoration failed: {e}")

    def _run_single_task(self, img_array: np.ndarray, task: str, save_prompts: bool) -> Tuple[np.ndarray, Dict]:
        """Run a single restoration task"""

        # Create temporary file for inference
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
            tmp_path = tmp_file.name
            cv2.imwrite(tmp_path, img_array)

        try:
            # Change to DocRes directory for inference to work properly
            original_cwd = os.getcwd()
            os.chdir(str(docres_dir))

            # Set global DEVICE variable that DocRes inference expects
            import inference  # Import the inference module to set its global DEVICE
            inference.DEVICE = self.device

            try:
                # Run inference
                prompt1, prompt2, prompt3, restored = inference_one_im(self._model, tmp_path, task)
            finally:
                # Always restore original working directory
                os.chdir(original_cwd)

            metadata = {
                'task': task,
                'device': str(self.device)
            }

            if save_prompts:
                metadata['prompts'] = {
                    'prompt1': prompt1,
                    'prompt2': prompt2, 
                    'prompt3': prompt3
                }

            return restored, metadata

        finally:
            # Clean up temporary file with retry for Windows
            try:
                # Wait a bit for file handles to be released
                time.sleep(0.1)
                os.unlink(tmp_path)
            except PermissionError:
                # If still locked, try again after a longer wait
                time.sleep(1)
                try:
                    os.unlink(tmp_path)
                except PermissionError:
                    # If still failing, just leave it - it will be cleaned up by the OS
                    pass

    def _run_end2end_pipeline(self, img_array: np.ndarray, save_prompts: bool) -> Tuple[np.ndarray, Dict]:
        """Run the end2end pipeline: dewarping → deshadowing → appearance"""

        intermediate_steps = {}

        # Change to DocRes directory for inference to work properly
        original_cwd = os.getcwd()
        os.chdir(str(docres_dir))

        # Set global DEVICE variable that DocRes inference expects
        import inference  # Import the inference module to set its global DEVICE
        inference.DEVICE = self.device

        try:
            with tempfile.TemporaryDirectory() as tmp_dir:
                # Step 1: Dewarping
                step1_path = os.path.join(tmp_dir, "step1.jpg")
                cv2.imwrite(step1_path, img_array)

                prompt1, prompt2, prompt3, dewarped = inference_one_im(self._model, step1_path, "dewarping")
                intermediate_steps['dewarped'] = dewarped

                # Step 2: Deshadowing
                step2_path = os.path.join(tmp_dir, "step2.jpg")
                cv2.imwrite(step2_path, dewarped)

                prompt1, prompt2, prompt3, deshadowed = inference_one_im(self._model, step2_path, "deshadowing")
                intermediate_steps['deshadowed'] = deshadowed

                # Step 3: Appearance
                step3_path = os.path.join(tmp_dir, "step3.jpg")
                cv2.imwrite(step3_path, deshadowed)

                prompt1, prompt2, prompt3, final = inference_one_im(self._model, step3_path, "appearance")

                metadata = {
                    'task': 'end2end',
                    'device': str(self.device),
                    'intermediate_steps': intermediate_steps
                }

                if save_prompts:
                    metadata['prompts'] = {
                        'prompt1': prompt1,
                        'prompt2': prompt2,
                        'prompt3': prompt3
                    }

                return final, metadata
        finally:
            # Always restore original working directory
            os.chdir(original_cwd)

    def batch_restore(
        self, 
        images: List[Union[str, np.ndarray]], 
        task: str = "appearance",
        save_prompts: bool = False
    ) -> List[Tuple[Optional[np.ndarray], Dict[str, Any]]]:
        """
        Restore multiple images in batch

        Args:
            images: List of image paths or numpy arrays
            task: Restoration task to perform
            save_prompts: Whether to save intermediate prompts

        Returns:
            List of (restored_image, metadata) tuples
        """
        results = []

        for i, image in enumerate(images):
            try:
                restored_img, metadata = self.restore_image(image, task, save_prompts)
                results.append((restored_img, metadata))
            except Exception as e:
                # Return None for failed images with error metadata
                error_metadata = {
                    'error': str(e),
                    'task': task,
                    'device': str(self.device),
                    'image_index': i
                }
                results.append((None, error_metadata))

        return results

    def get_supported_tasks(self) -> List[str]:
        """Get list of supported restoration tasks"""
        return self.SUPPORTED_TASKS.copy()

    def is_available(self) -> bool:
        """Check if DocRes is available and properly configured"""
        return DOCRES_AVAILABLE and self._model is not None

    def restore_pdf(
        self, 
        pdf_path: str, 
        output_path: str | None = None,
        task: str = "appearance",
        dpi: int = 200
    ) -> str | None:
        """
        Restore an entire PDF document using DocRes

        Args:
            pdf_path: Path to the input PDF file
            output_path: Path for the enhanced PDF (if None, auto-generates)
            task: DocRes restoration task (default: "appearance")
            dpi: DPI for PDF rendering (default: 200)

        Returns:
            Path to the enhanced PDF or None if failed
        """
        try:
            from PIL import Image
            from doctra.utils.pdf_io import render_pdf_to_images

            # Generate output path if not provided
            if output_path is None:
                pdf_dir = os.path.dirname(pdf_path)
                pdf_name = os.path.splitext(os.path.basename(pdf_path))[0]
                output_path = os.path.join(pdf_dir, f"{pdf_name}_enhanced.pdf")

            print(f"🔄 Processing PDF with DocRes: {os.path.basename(pdf_path)}")

            # Render all pages to images
            pil_pages = [im for (im, _, _) in render_pdf_to_images(pdf_path, dpi=dpi)]

            if not pil_pages:
                print("❌ No pages found in PDF")
                return None

            # Process each page with DocRes
            enhanced_pages = []

            # Detect environment for progress bar
            is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules

            # Create progress bar for page processing
            if is_notebook:
                progress_bar = create_notebook_friendly_bar(
                    total=len(pil_pages), 
                    desc="Processing pages"
                )
            else:
                progress_bar = create_beautiful_progress_bar(
                    total=len(pil_pages), 
                    desc="Processing pages",
                    leave=True
                )

            with progress_bar:
                for i, page_img in enumerate(pil_pages):
                    try:
                        # Convert PIL to numpy array
                        img_array = np.array(page_img)

                        # Apply DocRes restoration
                        restored_img, _ = self.restore_image(img_array, task)

                        # Convert back to PIL Image
                        enhanced_page = Image.fromarray(restored_img)
                        enhanced_pages.append(enhanced_page)

                        progress_bar.set_description(f"✅ Page {i+1}/{len(pil_pages)} processed")
                        progress_bar.update(1)

                    except Exception as e:
                        print(f"  ⚠️ Page {i+1} processing failed: {e}, using original")
                        enhanced_pages.append(page_img)
                        progress_bar.set_description(f"⚠️ Page {i+1} failed, using original")
                        progress_bar.update(1)

            # Create enhanced PDF
            if enhanced_pages:
                enhanced_pages[0].save(
                    output_path,
                    "PDF",
                    resolution=100.0,
                    save_all=True,
                    append_images=enhanced_pages[1:] if len(enhanced_pages) > 1 else []
                )

                print(f"✅ Enhanced PDF saved: {output_path}")
                return output_path
            else:
                print("❌ No pages to save")
                return None

        except ImportError as e:
            print(f"❌ Required dependencies not available: {e}")
            print("Install with: pip install PyMuPDF")
            return None
        except Exception as e:
            print(f"❌ Error processing PDF with DocRes: {e}")
            return None

__init__(device=None, use_half_precision=True, model_path=None, mbd_path=None)

Initialize DocRes Engine

Parameters:

Name	Type	Description	Default
device	Optional[str]	Device to run on ('cuda', 'cpu', or None for auto-detect)	None
use_half_precision	bool	Whether to use half precision for inference	True
model_path	Optional[str]	Path to DocRes model checkpoint (optional, defaults to Hugging Face Hub)	None
mbd_path	Optional[str]	Path to MBD model checkpoint (optional, defaults to Hugging Face Hub)	None

Source code in doctra/engines/image_restoration/docres_engine.py

def __init__(
    self, 
    device: Optional[str] = None,
    use_half_precision: bool = True,
    model_path: Optional[str] = None,
    mbd_path: Optional[str] = None
):
    """
    Initialize DocRes Engine

    Args:
        device: Device to run on ('cuda', 'cpu', or None for auto-detect)
        use_half_precision: Whether to use half precision for inference
        model_path: Path to DocRes model checkpoint (optional, defaults to Hugging Face Hub)
        mbd_path: Path to MBD model checkpoint (optional, defaults to Hugging Face Hub)
    """
    if not DOCRES_AVAILABLE:
        raise ImportError(
            "DocRes is not available. Please install the missing dependencies:\n"
            "pip install scikit-image>=0.19.3\n\n"
            "The DocRes module is already included in this library, but requires "
            "scikit-image for image processing operations."
        )

    # Set device
    if device is None:
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    else:
        requested_device = torch.device(device)
        # Check if the requested device is available
        if requested_device.type == 'cuda' and not torch.cuda.is_available():
            print(f"Warning: CUDA requested but not available. Falling back to CPU.")
            self.device = torch.device('cpu')
        else:
            self.device = requested_device

    self.use_half_precision = use_half_precision

    # Get model paths (always from Hugging Face Hub)
    try:
        self.mbd_path, self.model_path = get_model_paths(
            use_huggingface=True,
            model_path=model_path,
            mbd_path=mbd_path
        )
    except Exception as e:
        raise RuntimeError(f"Failed to get model paths: {e}")

    # Verify model files exist
    if not os.path.exists(self.model_path):
        raise FileNotFoundError(
            f"DocRes model not found at {self.model_path}. "
            f"This may indicate a Hugging Face download failure. "
            f"Please check your internet connection and try again."
        )

    if not os.path.exists(self.mbd_path):
        raise FileNotFoundError(
            f"MBD model not found at {self.mbd_path}. "
            f"This may indicate a Hugging Face download failure. "
            f"Please check your internet connection and try again."
        )

    # Initialize model
    self._model = None
    self._initialize_model()

batch_restore(images, task='appearance', save_prompts=False)

Restore multiple images in batch

Parameters:

Name	Type	Description	Default
images	List[Union[str, ndarray]]	List of image paths or numpy arrays	required
task	str	Restoration task to perform	'appearance'
save_prompts	bool	Whether to save intermediate prompts	False

Returns:

Type	Description
List[Tuple[Optional[ndarray], Dict[str, Any]]]	List of (restored_image, metadata) tuples

Source code in doctra/engines/image_restoration/docres_engine.py

def batch_restore(
    self, 
    images: List[Union[str, np.ndarray]], 
    task: str = "appearance",
    save_prompts: bool = False
) -> List[Tuple[Optional[np.ndarray], Dict[str, Any]]]:
    """
    Restore multiple images in batch

    Args:
        images: List of image paths or numpy arrays
        task: Restoration task to perform
        save_prompts: Whether to save intermediate prompts

    Returns:
        List of (restored_image, metadata) tuples
    """
    results = []

    for i, image in enumerate(images):
        try:
            restored_img, metadata = self.restore_image(image, task, save_prompts)
            results.append((restored_img, metadata))
        except Exception as e:
            # Return None for failed images with error metadata
            error_metadata = {
                'error': str(e),
                'task': task,
                'device': str(self.device),
                'image_index': i
            }
            results.append((None, error_metadata))

    return results

get_supported_tasks()

Get list of supported restoration tasks

Source code in doctra/engines/image_restoration/docres_engine.py

def get_supported_tasks(self) -> List[str]:
    """Get list of supported restoration tasks"""
    return self.SUPPORTED_TASKS.copy()

is_available()

Check if DocRes is available and properly configured

Source code in doctra/engines/image_restoration/docres_engine.py

def is_available(self) -> bool:
    """Check if DocRes is available and properly configured"""
    return DOCRES_AVAILABLE and self._model is not None

restore_image(image, task='appearance', save_prompts=False)

Restore a single image using DocRes

Parameters:

Name	Type	Description	Default
image	Union[str, ndarray]	Path to image file or numpy array	required
task	str	Restoration task to perform	'appearance'
save_prompts	bool	Whether to save intermediate prompts	False

Returns:

Type	Description
Tuple[ndarray, Dict[str, Any]]	Tuple of (restored_image, metadata)

Source code in doctra/engines/image_restoration/docres_engine.py

def restore_image(
    self, 
    image: Union[str, np.ndarray], 
    task: str = "appearance",
    save_prompts: bool = False
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Restore a single image using DocRes

    Args:
        image: Path to image file or numpy array
        task: Restoration task to perform
        save_prompts: Whether to save intermediate prompts

    Returns:
        Tuple of (restored_image, metadata)
    """
    if task not in self.SUPPORTED_TASKS:
        raise ValueError(f"Unsupported task: {task}. Supported tasks: {self.SUPPORTED_TASKS}")

    # Load image if path provided
    if isinstance(image, str):
        if not os.path.exists(image):
            raise FileNotFoundError(f"Image not found: {image}")
        img_array = cv2.imread(image)
        if img_array is None:
            raise ValueError(f"Could not load image: {image}")
    else:
        img_array = image.copy()

    original_shape = img_array.shape

    try:
        # Handle end2end pipeline
        if task == "end2end":
            return self._run_end2end_pipeline(img_array, save_prompts)

        # Run single task
        restored_img, metadata = self._run_single_task(img_array, task, save_prompts)

        metadata.update({
            'original_shape': original_shape,
            'restored_shape': restored_img.shape,
            'task': task,
            'device': str(self.device)
        })

        return restored_img, metadata

    except Exception as e:
        raise RuntimeError(f"Image restoration failed: {e}")

restore_pdf(pdf_path, output_path=None, task='appearance', dpi=200)

Restore an entire PDF document using DocRes

Parameters:

Name	Type	Description	Default
pdf_path	str	Path to the input PDF file	required
output_path	str | None	Path for the enhanced PDF (if None, auto-generates)	None
task	str	DocRes restoration task (default: "appearance")	'appearance'
dpi	int	DPI for PDF rendering (default: 200)	200

Returns:

Type	Description
str | None	Path to the enhanced PDF or None if failed

Source code in doctra/engines/image_restoration/docres_engine.py

def restore_pdf(
    self, 
    pdf_path: str, 
    output_path: str | None = None,
    task: str = "appearance",
    dpi: int = 200
) -> str | None:
    """
    Restore an entire PDF document using DocRes

    Args:
        pdf_path: Path to the input PDF file
        output_path: Path for the enhanced PDF (if None, auto-generates)
        task: DocRes restoration task (default: "appearance")
        dpi: DPI for PDF rendering (default: 200)

    Returns:
        Path to the enhanced PDF or None if failed
    """
    try:
        from PIL import Image
        from doctra.utils.pdf_io import render_pdf_to_images

        # Generate output path if not provided
        if output_path is None:
            pdf_dir = os.path.dirname(pdf_path)
            pdf_name = os.path.splitext(os.path.basename(pdf_path))[0]
            output_path = os.path.join(pdf_dir, f"{pdf_name}_enhanced.pdf")

        print(f"🔄 Processing PDF with DocRes: {os.path.basename(pdf_path)}")

        # Render all pages to images
        pil_pages = [im for (im, _, _) in render_pdf_to_images(pdf_path, dpi=dpi)]

        if not pil_pages:
            print("❌ No pages found in PDF")
            return None

        # Process each page with DocRes
        enhanced_pages = []

        # Detect environment for progress bar
        is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules

        # Create progress bar for page processing
        if is_notebook:
            progress_bar = create_notebook_friendly_bar(
                total=len(pil_pages), 
                desc="Processing pages"
            )
        else:
            progress_bar = create_beautiful_progress_bar(
                total=len(pil_pages), 
                desc="Processing pages",
                leave=True
            )

        with progress_bar:
            for i, page_img in enumerate(pil_pages):
                try:
                    # Convert PIL to numpy array
                    img_array = np.array(page_img)

                    # Apply DocRes restoration
                    restored_img, _ = self.restore_image(img_array, task)

                    # Convert back to PIL Image
                    enhanced_page = Image.fromarray(restored_img)
                    enhanced_pages.append(enhanced_page)

                    progress_bar.set_description(f"✅ Page {i+1}/{len(pil_pages)} processed")
                    progress_bar.update(1)

                except Exception as e:
                    print(f"  ⚠️ Page {i+1} processing failed: {e}, using original")
                    enhanced_pages.append(page_img)
                    progress_bar.set_description(f"⚠️ Page {i+1} failed, using original")
                    progress_bar.update(1)

        # Create enhanced PDF
        if enhanced_pages:
            enhanced_pages[0].save(
                output_path,
                "PDF",
                resolution=100.0,
                save_all=True,
                append_images=enhanced_pages[1:] if len(enhanced_pages) > 1 else []
            )

            print(f"✅ Enhanced PDF saved: {output_path}")
            return output_path
        else:
            print("❌ No pages to save")
            return None

    except ImportError as e:
        print(f"❌ Required dependencies not available: {e}")
        print("Install with: pip install PyMuPDF")
        return None
    except Exception as e:
        print(f"❌ Error processing PDF with DocRes: {e}")
        return None

---

Quick Reference

DocResEngine

from doctra import DocResEngine

# Initialize engine
engine = DocResEngine(
    device: str = None,  # "cuda", "cpu", or None for auto-detect
    use_half_precision: bool = False,
    model_path: str = None,
    mbd_path: str = None
)

# Restore single image
restored_img, metadata = engine.restore_image(
    image: Union[str, np.ndarray, PIL.Image.Image],
    task: str = "appearance"
)

# Restore PDF
output_path = engine.restore_pdf(
    pdf_path: str,
    output_path: str = None,
    task: str = "appearance",
    dpi: int = 200
)

Parameter Reference

Initialization Parameters

Parameter	Type	Default	Description
device	str	None	Processing device: "cuda", "cpu", or None (auto-detect)
use_half_precision	bool	False	Use FP16 for faster GPU processing
model_path	str	None	Custom path to restoration model
mbd_path	str	None	Custom path to MBD model

Restoration Tasks

Task	Description	Use Case
"appearance"	General appearance enhancement	Most documents (default)
"dewarping"	Correct perspective distortion	Scanned with perspective issues
"deshadowing"	Remove shadows and lighting artifacts	Poor lighting conditions
"deblurring"	Reduce blur and improve sharpness	Motion blur, focus issues
"binarization"	Convert to black and white	Clean text extraction
"end2end"	Complete restoration pipeline	Severely degraded documents

Methods

restore_image()

Restore a single image.

Parameters:

• image (str | np.ndarray | PIL.Image.Image): Input image (path, numpy array, or PIL Image)
• task (str): Restoration task to perform

Returns:

• restored_img (PIL.Image.Image): Restored image
• metadata (dict): Processing metadata including task, device, and timing

Example:

from doctra import DocResEngine

engine = DocResEngine(device="cuda")
restored, meta = engine.restore_image("blurry.jpg", task="deblurring")

print(f"Task: {meta['task']}")
print(f"Device: {meta['device']}")
print(f"Time: {meta['processing_time']:.2f}s")

# Save restored image
restored.save("restored.jpg")

restore_pdf()

Restore all pages in a PDF document.

Parameters:

• pdf_path (str): Path to input PDF
• output_path (str, optional): Path for output PDF (auto-generated if None)
• task (str): Restoration task to perform
• dpi (int): Resolution for processing

Returns:

• output_path (str): Path to the restored PDF

Example:

from doctra import DocResEngine

engine = DocResEngine(device="cuda")
restored_pdf = engine.restore_pdf(
    pdf_path="low_quality.pdf",
    output_path="enhanced.pdf",
    task="appearance",
    dpi=300
)

print(f"Restored PDF saved to: {restored_pdf}")

Device Selection

Auto-Detection

# Automatically uses GPU if available, otherwise CPU
engine = DocResEngine()

Explicit GPU

# Force GPU usage (will error if CUDA not available)
engine = DocResEngine(device="cuda")

Explicit CPU

# Force CPU usage (slower but always available)
engine = DocResEngine(device="cpu")

Check Device

import torch

print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

Performance Optimization

Half Precision

Use FP16 for ~2x speed on modern GPUs:

engine = DocResEngine(
    device="cuda",
    use_half_precision=True  # Faster, minimal quality loss
)

Requirements: - NVIDIA GPU with compute capability 7.0+ (Volta or newer) - Examples: RTX 20xx, RTX 30xx, RTX 40xx, A100, V100

Batch Processing

Process multiple images efficiently:

from doctra import DocResEngine

engine = DocResEngine(device="cuda")

# Process image list
images = ["doc1.jpg", "doc2.jpg", "doc3.jpg"]
restored_images = []

for img_path in images:
    restored, _ = engine.restore_image(img_path, task="appearance")
    restored_images.append(restored)
    restored.save(f"restored_{img_path}")

DPI Considerations

DPI	Quality	Speed	Memory	Best For
100	Low	Fast	Low	Quick previews
150	Medium	Medium	Medium	General use
200	Good	Slow	Medium	Default setting
300	High	Very Slow	High	High-quality scans

Metadata

The restore_image() method returns metadata:

restored, metadata = engine.restore_image("doc.jpg", "appearance")

print(metadata)
# {
#     'task': 'appearance',
#     'device': 'cuda',
#     'processing_time': 1.23,
#     'input_size': (1920, 1080),
#     'output_size': (1920, 1080)
# }

Error Handling

from doctra import DocResEngine

engine = DocResEngine(device="cuda")

try:
    restored, meta = engine.restore_image("document.jpg", "appearance")
except FileNotFoundError:
    print("Image not found")
except RuntimeError as e:
    print(f"CUDA error: {e}")
    # Fall back to CPU
    engine = DocResEngine(device="cpu")
    restored, meta = engine.restore_image("document.jpg", "appearance")
except Exception as e:
    print(f"Unexpected error: {e}")

Integration with Parsers

DocResEngine is integrated into EnhancedPDFParser:

from doctra import EnhancedPDFParser

# This internally uses DocResEngine
parser = EnhancedPDFParser(
    use_image_restoration=True,
    restoration_task="appearance",
    restoration_device="cuda"
)

parser.parse("document.pdf")

For standalone restoration:

from doctra import DocResEngine

# Step 1: Restore PDF
engine = DocResEngine(device="cuda")
enhanced_pdf = engine.restore_pdf(
    pdf_path="low_quality.pdf",
    output_path="enhanced.pdf",
    task="appearance"
)

# Step 2: Parse enhanced PDF
from doctra import StructuredPDFParser

parser = StructuredPDFParser()
parser.parse(enhanced_pdf)

Examples

Example 1: Dewarp Scanned Document

from doctra import DocResEngine

engine = DocResEngine(device="cuda")

# Fix perspective distortion
restored, meta = engine.restore_image(
    "scanned_with_distortion.jpg",
    task="dewarping"
)

restored.save("dewarped.jpg")
print(f"Processed in {meta['processing_time']:.2f}s")

Example 2: Remove Shadows

from doctra import DocResEngine

engine = DocResEngine(device="cuda")

# Remove shadow artifacts
restored, meta = engine.restore_image(
    "document_with_shadows.jpg",
    task="deshadowing"
)

restored.save("no_shadows.jpg")

Example 3: Batch PDF Restoration

import os
from doctra import DocResEngine

engine = DocResEngine(device="cuda", use_half_precision=True)

pdf_dir = "input_pdfs"
output_dir = "restored_pdfs"
os.makedirs(output_dir, exist_ok=True)

for filename in os.listdir(pdf_dir):
    if filename.endswith(".pdf"):
        input_path = os.path.join(pdf_dir, filename)
        output_path = os.path.join(output_dir, f"restored_{filename}")

        print(f"Processing {filename}...")
        engine.restore_pdf(
            pdf_path=input_path,
            output_path=output_path,
            task="appearance",
            dpi=200
        )

See Also

• Enhanced Parser - Using restoration with parsing
• Core Concepts - Understanding image restoration
• Examples - Advanced usage patterns