doctra 0.3.3__py3-none-any.whl → 0.4.1__py3-none-any.whl

doctra/exporters/excel_writer.py

@@ -5,6 +5,7 @@ from typing import Dict, Any, List, Set
 import pandas as pd  # pip install pandas openpyxl
 from openpyxl.styles import PatternFill, Font, Alignment
 from openpyxl.utils import get_column_letter
+from openpyxl.worksheet.hyperlink import Hyperlink
 
 _INVALID_SHEET_CHARS = r'[:\\/*?\[\]]'  # Excel-invalid characters
 _MAX_SHEET_LEN = 31
@@ -85,6 +86,61 @@ def _autosize_columns(ws, df: pd.DataFrame) -> None:
         ws.column_dimensions[get_column_letter(i)].width = min(max(10, max_len + 2), 60)
 
 
+def _style_summary_sheet(ws, df: pd.DataFrame, sheet_mapping: dict = None) -> None:
+    """
+    Apply special styling to the summary sheet with text wrapping for descriptions.
+    Add hyperlinks to table titles that link to their corresponding sheets.
+    
+    :param ws: OpenPyXL worksheet object to style
+    :param df: Pandas DataFrame containing the summary data
+    :param sheet_mapping: Dictionary mapping table titles to their sheet names
+    :return: None
+    """
+    # Style header row
+    _style_header(ws, ncols=df.shape[1])
+    
+    # Apply text wrapping to all data cells
+    wrap_alignment = Alignment(wrap_text=True, vertical="top")
+    
+    # Apply wrapping to all data rows (skip header row)
+    for row_idx in range(2, len(df) + 2):  # Start from row 2 (after header)
+        for col_idx in range(1, df.shape[1] + 1):
+            cell = ws.cell(row=row_idx, column=col_idx)
+            cell.alignment = wrap_alignment
+            
+            # Add hyperlink to table title column (column A)
+            if col_idx == 1 and sheet_mapping:  # Table Title column
+                table_title = cell.value
+                if table_title and table_title in sheet_mapping:
+                    sheet_name = sheet_mapping[table_title]
+                    
+                    # Create hyperlink to the sheet using proper Excel format
+                    # Escape sheet name if it contains spaces or special characters
+                    if ' ' in sheet_name or any(char in sheet_name for char in ['[', ']', '*', '?', ':', '\\', '/']):
+                        hyperlink_ref = f"#'{sheet_name}'!A1"
+                    else:
+                        hyperlink_ref = f"#{sheet_name}!A1"
+                    
+                    # Use Hyperlink class with proper parameters
+                    cell.hyperlink = Hyperlink(ref=hyperlink_ref, target=hyperlink_ref)
+                    # Style the hyperlink
+                    cell.font = Font(color="0000FF", underline="single")
+    
+    # Set specific column widths for summary sheet
+    # Table Title column - narrower
+    ws.column_dimensions['A'].width = 30
+    # Description column - wider to accommodate wrapped text
+    ws.column_dimensions['B'].width = 60
+    # Page column - narrow for page numbers
+    ws.column_dimensions['C'].width = 10
+    # Type column - narrow for Table/Chart
+    ws.column_dimensions['D'].width = 12
+    
+    # Set row heights to accommodate wrapped text
+    for row_idx in range(2, len(df) + 2):
+        ws.row_dimensions[row_idx].height = 60  # Allow for multiple lines
+
+
 def _normalize_data(headers: List[str], rows: List[List]) -> tuple[List[str], List[List]]:
     """
     Normalize headers and rows to ensure consistent dimensions.
@@ -159,6 +215,31 @@ def write_structured_excel(excel_path: str, items: List[Dict[str, Any]]) -> str
     taken: Set[str] = set()
 
     with pd.ExcelWriter(excel_path, engine="openpyxl", mode="w") as writer:
+        # Create summary sheet first
+        summary_data = []
+        sheet_mapping = {}  # Map table titles to their sheet names
+        
+        for item in valid_items:
+            title = item.get("title") or "Untitled"
+            description = item.get("description") or "No description available"
+            page_number = item.get("page", "Unknown")
+            item_type = item.get("type", "Table")  # Default to "Table" if not specified
+            
+            
+            summary_data.append({
+                "Table Title": title,
+                "Description": description,
+                "Page": page_number,
+                "Type": item_type
+            })
+        
+        # Create summary sheet first (but without hyperlinks initially)
+        if summary_data:
+            summary_df = pd.DataFrame(summary_data)
+            summary_df.to_excel(writer, sheet_name="Table Summary", index=False)
+            taken.add("Table Summary")
+
+        # Process individual table sheets to build sheet mapping
         for item in valid_items:
             try:
                 title = item.get("title") or "Untitled"
@@ -166,6 +247,9 @@ def write_structured_excel(excel_path: str, items: List[Dict[str, Any]]) -> str
                 rows = item.get("rows") or []
 
                 sheet_name = _safe_sheet_name(title, taken)
+                
+                # Add to sheet mapping for hyperlinks
+                sheet_mapping[title] = sheet_name
 
                 # Normalize data to handle mismatched dimensions
                 normalized_headers, normalized_rows = _normalize_data(headers, rows)
@@ -194,4 +278,9 @@ def write_structured_excel(excel_path: str, items: List[Dict[str, Any]]) -> str
                 print(f"Error processing item '{item.get('title', 'Unknown')}': {e}")
                 continue
 
+        # Now add hyperlinks to the summary sheet (after all sheets are created)
+        if summary_data and sheet_mapping:
+            summary_ws = writer.sheets["Table Summary"]
+            _style_summary_sheet(summary_ws, summary_df, sheet_mapping)
+
     return excel_path

doctra/parsers/enhanced_pdf_parser.py

@@ -0,0 +1,374 @@
+"""
+Enhanced PDF Parser with Image Restoration
+
+This module provides an enhanced PDF parser that combines the structured parsing
+capabilities with DocRes image restoration for improved document processing.
+"""
+
+from __future__ import annotations
+import os
+import sys
+import numpy as np
+from typing import List, Dict, Any, Optional, Union
+from contextlib import ExitStack
+from PIL import Image
+from tqdm import tqdm
+
+from doctra.parsers.structured_pdf_parser import StructuredPDFParser
+from doctra.engines.image_restoration import DocResEngine
+from doctra.utils.pdf_io import render_pdf_to_images
+from doctra.utils.constants import IMAGE_SUBDIRS, EXCLUDE_LABELS
+from doctra.utils.file_ops import ensure_output_dirs
+from doctra.utils.progress import create_beautiful_progress_bar, create_notebook_friendly_bar
+from doctra.parsers.layout_order import reading_order_key
+from doctra.utils.ocr_utils import ocr_box_text
+from doctra.exporters.image_saver import save_box_image
+from doctra.exporters.markdown_writer import write_markdown
+from doctra.exporters.html_writer import write_html, write_structured_html
+from doctra.exporters.excel_writer import write_structured_excel
+from doctra.utils.structured_utils import to_structured_dict
+from doctra.exporters.markdown_table import render_markdown_table
+
+
+class EnhancedPDFParser(StructuredPDFParser):
+    """
+    Enhanced PDF Parser with Image Restoration capabilities.
+    
+    Extends the StructuredPDFParser with DocRes image restoration to improve
+    document quality before processing. This is particularly useful for:
+    - Scanned documents with shadows or distortion
+    - Low-quality PDFs that need enhancement
+    - Documents with perspective issues
+    
+    :param use_image_restoration: Whether to apply DocRes image restoration (default: True)
+    :param restoration_task: DocRes task to use ("dewarping", "deshadowing", "appearance", "deblurring", "binarization", "end2end", default: "appearance")
+    :param restoration_device: Device for DocRes processing ("cuda", "cpu", or None for auto-detect, default: None)
+    :param restoration_dpi: DPI for restoration processing (default: 200)
+    :param use_vlm: Whether to use VLM for structured data extraction (default: False)
+    :param vlm_provider: VLM provider to use ("gemini", "openai", "anthropic", or "openrouter", default: "gemini")
+    :param vlm_model: Model name to use (defaults to provider-specific defaults)
+    :param vlm_api_key: API key for VLM provider (required if use_vlm is True)
+    :param layout_model_name: Layout detection model name (default: "PP-DocLayout_plus-L")
+    :param dpi: DPI for PDF rendering (default: 200)
+    :param min_score: Minimum confidence score for layout detection (default: 0.0)
+    :param ocr_lang: OCR language code (default: "eng")
+    :param ocr_psm: Tesseract page segmentation mode (default: 4)
+    :param ocr_oem: Tesseract OCR engine mode (default: 3)
+    :param ocr_extra_config: Additional Tesseract configuration (default: "")
+    :param box_separator: Separator between text boxes in output (default: "\n")
+    """
+
+    def __init__(
+        self,
+        *,
+        use_image_restoration: bool = True,
+        restoration_task: str = "appearance",
+        restoration_device: Optional[str] = None,
+        restoration_dpi: int = 200,
+        use_vlm: bool = False,
+        vlm_provider: str = "gemini",
+        vlm_model: str | None = None,
+        vlm_api_key: str | None = None,
+        layout_model_name: str = "PP-DocLayout_plus-L",
+        dpi: int = 200,
+        min_score: float = 0.0,
+        ocr_lang: str = "eng",
+        ocr_psm: int = 4,
+        ocr_oem: int = 3,
+        ocr_extra_config: str = "",
+        box_separator: str = "\n",
+    ):
+        """
+        Initialize the Enhanced PDF Parser with image restoration capabilities.
+        """
+        # Initialize parent class
+        super().__init__(
+            use_vlm=use_vlm,
+            vlm_provider=vlm_provider,
+            vlm_model=vlm_model,
+            vlm_api_key=vlm_api_key,
+            layout_model_name=layout_model_name,
+            dpi=dpi,
+            min_score=min_score,
+            ocr_lang=ocr_lang,
+            ocr_psm=ocr_psm,
+            ocr_oem=ocr_oem,
+            ocr_extra_config=ocr_extra_config,
+            box_separator=box_separator,
+        )
+        
+        # Image restoration settings
+        self.use_image_restoration = use_image_restoration
+        self.restoration_task = restoration_task
+        self.restoration_device = restoration_device
+        self.restoration_dpi = restoration_dpi
+        
+        # Initialize DocRes engine if needed
+        self.docres_engine = None
+        if self.use_image_restoration:
+            try:
+                self.docres_engine = DocResEngine(
+                    device=restoration_device,
+                    use_half_precision=True
+                )
+                print(f"✅ DocRes engine initialized with task: {restoration_task}")
+            except Exception as e:
+                print(f"⚠️ DocRes initialization failed: {e}")
+                print("   Continuing without image restoration...")
+                self.use_image_restoration = False
+                self.docres_engine = None
+
+    def parse(self, pdf_path: str, enhanced_output_dir: str = None) -> None:
+        """
+        Parse a PDF document with optional image restoration.
+        
+        :param pdf_path: Path to the input PDF file
+        :param enhanced_output_dir: Directory for enhanced images (if None, uses default)
+        :return: None
+        """
+        pdf_filename = os.path.splitext(os.path.basename(pdf_path))[0]
+        
+        # Set up output directories
+        if enhanced_output_dir is None:
+            out_dir = f"outputs/{pdf_filename}/enhanced_parse"
+        else:
+            out_dir = enhanced_output_dir
+            
+        os.makedirs(out_dir, exist_ok=True)
+        ensure_output_dirs(out_dir, IMAGE_SUBDIRS)
+        
+        # Process PDF pages with optional restoration
+        if self.use_image_restoration and self.docres_engine:
+            print(f"🔄 Processing PDF with image restoration: {os.path.basename(pdf_path)}")
+            enhanced_pages = self._process_pages_with_restoration(pdf_path, out_dir)
+        else:
+            print(f"🔄 Processing PDF without image restoration: {os.path.basename(pdf_path)}")
+            enhanced_pages = [im for (im, _, _) in render_pdf_to_images(pdf_path, dpi=self.dpi)]
+        
+        # Run layout detection on enhanced pages
+        print("🔍 Running layout detection on enhanced pages...")
+        pages = self.layout_engine.predict_pdf(
+            pdf_path, batch_size=1, layout_nms=True, dpi=self.dpi, min_score=self.min_score
+        )
+        
+        # Use enhanced pages for processing
+        pil_pages = enhanced_pages
+        
+        # Continue with standard parsing logic
+        self._process_parsing_logic(pages, pil_pages, out_dir, pdf_filename, pdf_path)
+
+    def _process_pages_with_restoration(self, pdf_path: str, out_dir: str) -> List[Image.Image]:
+        """
+        Process PDF pages with DocRes image restoration.
+        
+        :param pdf_path: Path to the input PDF file
+        :param out_dir: Output directory for enhanced images
+        :return: List of enhanced PIL images
+        """
+        # Render original pages
+        original_pages = [im for (im, _, _) in render_pdf_to_images(pdf_path, dpi=self.restoration_dpi)]
+        
+        if not original_pages:
+            print("❌ No pages found in PDF")
+            return []
+        
+        # Create progress bar
+        is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules
+        if is_notebook:
+            progress_bar = create_notebook_friendly_bar(
+                total=len(original_pages), 
+                desc=f"🔄 DocRes {self.restoration_task}"
+            )
+        else:
+            progress_bar = create_beautiful_progress_bar(
+                total=len(original_pages), 
+                desc=f"🔄 DocRes {self.restoration_task}",
+                leave=True
+            )
+        
+        enhanced_pages = []
+        enhanced_dir = os.path.join(out_dir, "enhanced_pages")
+        os.makedirs(enhanced_dir, exist_ok=True)
+        
+        try:
+            with progress_bar:
+                for i, page_img in enumerate(original_pages):
+                    try:
+                        # Convert PIL to numpy array
+                        img_array = np.array(page_img)
+                        
+                        # Apply DocRes restoration
+                        restored_img, metadata = self.docres_engine.restore_image(
+                            img_array, 
+                            task=self.restoration_task
+                        )
+                        
+                        # Convert back to PIL Image
+                        enhanced_page = Image.fromarray(restored_img)
+                        enhanced_pages.append(enhanced_page)
+                        
+                        # Save enhanced page for reference
+                        enhanced_path = os.path.join(enhanced_dir, f"page_{i+1:03d}_enhanced.jpg")
+                        enhanced_page.save(enhanced_path, "JPEG", quality=95)
+                        
+                        progress_bar.set_description(f"✅ Page {i+1}/{len(original_pages)} enhanced")
+                        progress_bar.update(1)
+                        
+                    except Exception as e:
+                        print(f"  ⚠️ Page {i+1} restoration failed: {e}, using original")
+                        enhanced_pages.append(page_img)
+                        progress_bar.set_description(f"⚠️ Page {i+1} failed, using original")
+                        progress_bar.update(1)
+        
+        finally:
+            if hasattr(progress_bar, 'close'):
+                progress_bar.close()
+        
+        print(f"✅ Image restoration completed. Enhanced pages saved to: {enhanced_dir}")
+        return enhanced_pages
+
+    def _process_parsing_logic(self, pages, pil_pages, out_dir, pdf_filename, pdf_path):
+        """
+        Process the parsing logic with enhanced pages.
+        This is extracted from the parent class to allow customization.
+        """
+        
+        fig_count = sum(sum(1 for b in p.boxes if b.label == "figure") for p in pages)
+        chart_count = sum(sum(1 for b in p.boxes if b.label == "chart") for p in pages)
+        table_count = sum(sum(1 for b in p.boxes if b.label == "table") for p in pages)
+
+        md_lines: List[str] = ["# Enhanced Document Content\n"]
+        structured_items: List[Dict[str, Any]] = []
+
+        charts_desc = "Charts (VLM → table)" if self.use_vlm else "Charts (cropped)"
+        tables_desc = "Tables (VLM → table)" if self.use_vlm else "Tables (cropped)"
+        figures_desc = "Figures (cropped)"
+
+        with ExitStack() as stack:
+            is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules
+            if is_notebook:
+                charts_bar = stack.enter_context(
+                    create_notebook_friendly_bar(total=chart_count, desc=charts_desc)) if chart_count else None
+                tables_bar = stack.enter_context(
+                    create_notebook_friendly_bar(total=table_count, desc=tables_desc)) if table_count else None
+                figures_bar = stack.enter_context(
+                    create_notebook_friendly_bar(total=fig_count, desc=figures_desc)) if fig_count else None
+            else:
+                charts_bar = stack.enter_context(
+                    create_beautiful_progress_bar(total=chart_count, desc=charts_desc, leave=True)) if chart_count else None
+                tables_bar = stack.enter_context(
+                    create_beautiful_progress_bar(total=table_count, desc=tables_desc, leave=True)) if table_count else None
+                figures_bar = stack.enter_context(
+                    create_beautiful_progress_bar(total=fig_count, desc=figures_desc, leave=True)) if fig_count else None
+
+            for p in pages:
+                page_num = p.page_index
+                page_img: Image.Image = pil_pages[page_num - 1]
+                md_lines.append(f"\n## Page {page_num}\n")
+
+                for i, box in enumerate(sorted(p.boxes, key=reading_order_key), start=1):
+                    if box.label in EXCLUDE_LABELS:
+                        img_path = save_box_image(page_img, box, out_dir, page_num, i, IMAGE_SUBDIRS)
+                        abs_img_path = os.path.abspath(img_path)
+                        rel = os.path.relpath(abs_img_path, out_dir)
+
+                        if box.label == "figure":
+                            md_lines.append(f"![Figure — page {page_num}]({rel})\n")
+                            if figures_bar: figures_bar.update(1)
+
+                        elif box.label == "chart":
+                            if self.use_vlm and self.vlm:
+                                wrote_table = False
+                                try:
+                                    chart = self.vlm.extract_chart(abs_img_path)
+                                    item = to_structured_dict(chart)
+                                    if item:
+                                        # Add page and type information to structured item
+                                        item["page"] = page_num
+                                        item["type"] = "Chart"
+                                        structured_items.append(item)
+                                        md_lines.append(
+                                            render_markdown_table(item.get("headers"), item.get("rows"),
+                                                                  title=item.get("title"))
+                                        )
+                                        wrote_table = True
+                                except Exception as e:
+                                    pass
+                                if not wrote_table:
+                                    md_lines.append(f"![Chart — page {page_num}]({rel})\n")
+                            else:
+                                md_lines.append(f"![Chart — page {page_num}]({rel})\n")
+                            if charts_bar: charts_bar.update(1)
+
+                        elif box.label == "table":
+                            if self.use_vlm and self.vlm:
+                                wrote_table = False
+                                try:
+                                    table = self.vlm.extract_table(abs_img_path)
+                                    item = to_structured_dict(table)
+                                    if item:
+                                        # Add page and type information to structured item
+                                        item["page"] = page_num
+                                        item["type"] = "Table"
+                                        structured_items.append(item)
+                                        md_lines.append(
+                                            render_markdown_table(item.get("headers"), item.get("rows"),
+                                                                  title=item.get("title"))
+                                        )
+                                        wrote_table = True
+                                except Exception as e:
+                                    pass
+                                if not wrote_table:
+                                    md_lines.append(f"![Table — page {page_num}]({rel})\n")
+                            else:
+                                md_lines.append(f"![Table — page {page_num}]({rel})\n")
+                            if tables_bar: tables_bar.update(1)
+                    else:
+                        text = ocr_box_text(self.ocr_engine, page_img, box)
+                        if text:
+                            md_lines.append(text)
+                            md_lines.append(self.box_separator if self.box_separator else "")
+
+        md_path = write_markdown(md_lines, out_dir)
+        html_path = write_html(md_lines, out_dir)
+        
+        excel_path = None
+        html_structured_path = None
+        if self.use_vlm and structured_items:
+            excel_path = os.path.join(out_dir, "tables.xlsx")
+            write_structured_excel(excel_path, structured_items)
+            html_structured_path = os.path.join(out_dir, "tables.html")
+            write_structured_html(html_structured_path, structured_items)
+
+        print(f"✅ Enhanced parsing completed successfully!")
+        print(f"📁 Output directory: {out_dir}")
+
+    def restore_pdf_only(self, pdf_path: str, output_path: str = None, task: str = None) -> str:
+        """
+        Apply DocRes restoration to a PDF without parsing.
+        
+        :param pdf_path: Path to the input PDF file
+        :param output_path: Path for the enhanced PDF (if None, auto-generates)
+        :param task: DocRes restoration task (if None, uses instance default)
+        :return: Path to the enhanced PDF or None if failed
+        """
+        if not self.use_image_restoration or not self.docres_engine:
+            raise RuntimeError("Image restoration is not enabled or DocRes engine is not available")
+        
+        task = task or self.restoration_task
+        return self.docres_engine.restore_pdf(pdf_path, output_path, task, self.restoration_dpi)
+
+    def get_restoration_info(self) -> Dict[str, Any]:
+        """
+        Get information about the current restoration configuration.
+        
+        :return: Dictionary with restoration settings and status
+        """
+        return {
+            'enabled': self.use_image_restoration,
+            'task': self.restoration_task,
+            'device': self.restoration_device,
+            'dpi': self.restoration_dpi,
+            'engine_available': self.docres_engine is not None,
+            'supported_tasks': self.docres_engine.get_supported_tasks() if self.docres_engine else []
+        }

doctra/parsers/structured_pdf_parser.py

@@ -163,6 +163,9 @@ class StructuredPDFParser:
                                     chart = self.vlm.extract_chart(abs_img_path)
                                     item = to_structured_dict(chart)
                                     if item:
+                                        # Add page and type information to structured item
+                                        item["page"] = page_num
+                                        item["type"] = "Chart"
                                         structured_items.append(item)
                                         md_lines.append(
                                             render_markdown_table(item.get("headers"), item.get("rows"),
@@ -184,6 +187,9 @@ class StructuredPDFParser:
                                     table = self.vlm.extract_table(abs_img_path)
                                     item = to_structured_dict(table)
                                     if item:
+                                        # Add page and type information to structured item
+                                        item["page"] = page_num
+                                        item["type"] = "Table"
                                         structured_items.append(item)
                                         md_lines.append(
                                             render_markdown_table(item.get("headers"), item.get("rows"),

doctra/parsers/table_chart_extractor.py

@@ -178,6 +178,9 @@ class ChartTablePDFParser:
                                 extracted_chart = self.vlm.extract_chart(chart_path)
                                 structured_item = to_structured_dict(extracted_chart)
                                 if structured_item:
+                                    # Add page and type information to structured item
+                                    structured_item["page"] = page_num
+                                    structured_item["type"] = "Chart"
                                     structured_items.append(structured_item)
                                     vlm_items.append({
                                         "kind": "chart",
@@ -221,6 +224,9 @@ class ChartTablePDFParser:
                                 extracted_table = self.vlm.extract_table(table_path)
                                 structured_item = to_structured_dict(extracted_table)
                                 if structured_item:
+                                    # Add page and type information to structured item
+                                    structured_item["page"] = page_num
+                                    structured_item["type"] = "Table"
                                     structured_items.append(structured_item)
                                     vlm_items.append({
                                         "kind": "table",

doctra/third_party/docres/data/MBD/MBD.py

@@ -0,0 +1,110 @@
+import cv2
+import numpy as np
+import MBD_utils
+import torch
+import torch.nn.functional as F
+
+
+def mask_base_dewarper(image,mask):
+    '''
+    input:
+        image -> ndarray HxWx3 uint8
+        mask -> ndarray HxW uint8
+    return
+        dewarped -> ndarray HxWx3 uint8
+        grid (optional) -> ndarray HxWx2 -1~1
+    '''
+
+    ## get contours
+    # _, contours, hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)  ## cv2.__version__ == 3.x
+    contours,hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,method=cv2.CHAIN_APPROX_SIMPLE)  ## cv2.__version__ == 4.x
+
+    ## get biggest contours and four corners based on Douglas-Peucker algorithm
+    four_corners, maxArea, contour= MBD_utils.DP_algorithm(contours)
+    four_corners = MBD_utils.reorder(four_corners)
+
+    ## reserve biggest contours and remove other noisy contours
+    new_mask = np.zeros_like(mask)
+    new_mask = cv2.drawContours(new_mask,[contour],-1,255,cv2.FILLED)
+
+    ## obtain middle points
+    # ratios = [0.25,0.5,0.75]  # ratios = [0.125,0.25,0.375,0.5,0.625,0.75,0.875]
+    ratios = [0.25,0.5,0.75]
+    # ratios = [0.0625,0.125,0.1875,0.25,0.3125,0.375,0.4475,0.5,0.5625,0.625,0.06875,0.75,0.8125,0.875,0.9375]
+    middle = MBD_utils.findMiddle(corners=four_corners,mask=new_mask,points=ratios)
+    
+    ## all points
+    source_points = np.concatenate((four_corners,middle),axis=0) ## all_point = four_corners(topleft,topright,bottom)+top+bottom+left+right
+
+    ## target points
+    h,w = image.shape[:2]
+    padding = 0
+    target_points = [[padding, padding],[w-padding, padding], [padding, h-padding],[w-padding, h-padding]]
+    for ratio in ratios:
+        target_points.append([int((w-2*padding)*ratio)+padding,padding])
+    for ratio in ratios:
+        target_points.append([int((w-2*padding)*ratio)+padding,h-padding])
+    for ratio in ratios:
+        target_points.append([padding,int((h-2*padding)*ratio)+padding])
+    for ratio in ratios:
+        target_points.append([w-padding,int((h-2*padding)*ratio)+padding])
+
+    ## dewarp base on cv2
+    # pts1 = np.float32(source_points)
+    # pts2 = np.float32(target_points)
+    # tps = cv2.createThinPlateSplineShapeTransformer()
+    # matches = []
+    # N = pts1.shape[0]
+    # for i in range(0,N):
+    #     matches.append(cv2.DMatch(i,i,0))
+    # pts1 = pts1.reshape(1,-1,2)
+    # pts2 = pts2.reshape(1,-1,2)
+    # tps.estimateTransformation(pts2,pts1,matches)
+    # dewarped = tps.warpImage(image)
+
+    ## dewarp base on generated grid
+    source_points = source_points.reshape(-1,2)/np.array([image.shape[:2][::-1]]).reshape(1,2)
+    source_points = torch.from_numpy(source_points).float().cuda()
+    source_points = source_points.unsqueeze(0)
+    source_points = (source_points-0.5)*2
+    target_points = np.asarray(target_points).reshape(-1,2)/np.array([image.shape[:2][::-1]]).reshape(1,2)
+    target_points = torch.from_numpy(target_points).float()
+    target_points = (target_points-0.5)*2
+
+    model = MBD_utils.TPSGridGen(target_height=256,target_width=256,target_control_points=target_points)
+    model = model.cuda()
+    grid = model(source_points).view(-1,256,256,2).permute(0,3,1,2)
+    grid = F.interpolate(grid,(h,w),mode='bilinear').permute(0,2,3,1)
+    dewarped = MBD_utils.torch2cvimg(F.grid_sample(MBD_utils.cvimg2torch(image).cuda(),grid))[0]
+    return dewarped,grid[0].cpu().numpy()
+
+def mask_base_cropper(image,mask):
+    '''
+    input:
+        image -> ndarray HxWx3 uint8
+        mask -> ndarray HxW uint8
+    return
+        dewarped -> ndarray HxWx3 uint8
+        grid (optional) -> ndarray HxWx2 -1~1
+    '''
+
+    ## get contours
+    _, contours, hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)  ## cv2.__version__ == 3.x
+    # contours,hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,method=cv2.CHAIN_APPROX_SIMPLE)  ## cv2.__version__ == 4.x
+
+    ## get biggest contours and four corners based on Douglas-Peucker algorithm
+    four_corners, maxArea, contour= MBD_utils.DP_algorithm(contours)
+    four_corners = MBD_utils.reorder(four_corners)
+
+    ## reserve biggest contours and remove other noisy contours
+    new_mask = np.zeros_like(mask)
+    new_mask = cv2.drawContours(new_mask,[contour],-1,255,cv2.FILLED)
+
+    ## 最小外接矩形
+    rect = cv2.minAreaRect(contour) # 得到最小外接矩形的（中心(x,y), (宽,高), 旋转角度）
+    box = cv2.boxPoints(rect) # cv2.boxPoints(rect) for OpenCV 3.x 获取最小外接矩形的4个顶点坐标
+    box = np.int0(box)
+    box = box.reshape((4,1,2))    
+
+
+