Function process_element

demo_element.py:121–166

Process a single element image (text, table, formula) Args: image_path: Path to the element image model: HFModel model instance element_type: Type of element ('text', 'table', 'formula') save_dir: Directory to save results (default: same as input directory)

(image_path, model, element_type, save_dir=None)

def process_element(image_path, model, element_type, save_dir=None):
    """Process a single element image (text, table, formula)
    
    Args:
        image_path: Path to the element image
        model: HFModel model instance
        element_type: Type of element ('text', 'table', 'formula')
        save_dir: Directory to save results (default: same as input directory)
        
    Returns:
        Parsed content of the element and recognition results
    """
    # Load and prepare image
    pil_image = Image.open(image_path).convert("RGB")
    # pil_image = crop_margin(pil_image)
    
    # Select appropriate prompt based on element type
    if element_type == "table":
        prompt = "Parse the table in the image."
        label = "tab"
    elif element_type == "formula":
        prompt = "Read formula in the image."
        label = "equ"
    elif element_type == "code":
        prompt = "Read code in the image."
        label = "code"
    else:  # Default to text
        prompt = "Read text in the image."
        label = "para"
    
    # Process the element
    result = model.chat(prompt, pil_image)
    
    # Create recognition result in the same format as the document parser
    recognition_results = [
        {
            "label": label,
            "text": result.strip(),
        }
    ]
    
    # Save results if save_dir is provided
    save_outputs(recognition_results, pil_image, os.path.basename(image_path).split(".")[0], save_dir)
    print(f"Results saved to {save_dir}")
    
    return result, recognition_results


def main():