utils.py

import re
import os
import json
import random
from typing import Tuple, List, Dict
import numpy as np
from PIL import Image, ImageDraw, ImageFont

def clean_repeated_substrings(text):
    """Clean repeated substrings in text"""
    n = len(text)
    if n<8000:
        return text
    for length in range(2, n // 10 + 1):
        candidate = text[-length:] 
        count = 0
        i = n - length
        
        while i >= 0 and text[i:i + length] == candidate:
            count += 1
            i -= length

        if count >= 10:
            return text[:n - length * (count - 1)]  

    return text

def norm_formula_HYOCR(gt_form,pre_form):
    """
    NORMARLIZE for OmniDocBench.input:gt_form,pred_form(from quick_match_display_formula)
    # End-to-end document parsing bypasses layout constraints and category definitions, 
    # yet the HunyuanOCR model produces outputs with more flexible granularities. 
    # For certain mathematical expressions, HunyuanOCR may parse them using a mix of plain texts and embedded LaTeX formulas. 
    # This can lead to mismatches when evaluated with OmniDocBench’s quick-match mechanism. 
    # Therefore, we manually correct cases with incorrect matches (i.e., those with an edit-distance score greater than 0.5). 
    # Below is a piece of logic code used to handle some normalization cases.
    """
    def clean_gt_tail(gt: str):
        pattern = r'(\\quad+|\\qquad+)\s*\{?\(\s*\d+\s*\)\}?\s*$'
        return re.sub(pattern, '', gt).rstrip()
    
    pre_form = clean_gt_tail(pre_form)
    gt_form = gt_form.replace('\[', '').replace('\]', '').replace('，', ',').strip()
    pre_form = pre_form.replace('\[', '').replace('\]', '').replace('，', ',').strip()
    pre_form = pre_form.replace("\\text{ⓐ}", "\\textcircled{a}") \
        .replace("\\text{ⓑ}", "\\textcircled{b}") \
        .replace("\\text{ⓒ}", "\\textcircled{c}") \
        .replace("\\text{ⓓ}", "\\textcircled{d}")
    
    return gt_form, pre_form



def parse_coords(coord_str: str) -> Tuple[float, float]:
    """Parse coordinate string and return (x,y) tuple"""
    try:
        x, y = coord_str.strip('()').split(',')
        return (float(x), float(y))
    except:
        return (0, 0)

def denormalize_coordinates(coord: Tuple[float, float], image_width: int, image_height: int) -> Tuple[int, int]:
    """Denormalize coordinates from [0,1000] to image dimensions""" 
    x, y = coord
    denorm_x = int(x * image_width / 1000)
    denorm_y = int(y * image_height / 1000)
    return (denorm_x, denorm_y)

def process_spotting_response(response: str, image_width: int, image_height: int) -> str:
    """Process spotting task response and denormalize coordinates"""
    try:
        # Find all text and coordinate pairs using regex
        pattern = r'([^()]+)(\(\d+,\d+\),\(\d+,\d+\))'
        matches = re.finditer(pattern, response)
        
        new_response = response
        for match in matches:
            text = match.group(1).strip()
            coords = match.group(2)
            
            # Parse the two coordinate points 
            coord_pattern = r'\((\d+),(\d+)\)'
            coord_matches = re.findall(coord_pattern, coords)
            if len(coord_matches) == 2:
                start_coord = (float(coord_matches[0][0]), float(coord_matches[0][1]))
                end_coord = (float(coord_matches[1][0]), float(coord_matches[1][1]))
                
                # Denormalize coordinates
                denorm_start = denormalize_coordinates(start_coord, image_width, image_height)
                denorm_end = denormalize_coordinates(end_coord, image_width, image_height)
                
                # Build new coordinate string
                new_coords = f"({denorm_start[0]},{denorm_start[1]}),({denorm_end[0]},{denorm_end[1]})"
                
                # Replace coordinates in original response
                new_response = new_response.replace(coords, new_coords)
        
        return new_response
    
    except Exception as e:
        print(f"Error processing response: {str(e)}")
        return response

def draw_text_detection_boxes(image: Image, response: str) -> Image:
    """Draw text detection boxes on image"""
    image_width, image_height = image.size
    img_draw = image.copy()
    draw = ImageDraw.Draw(img_draw)
    
    # Create transparent overlay
    overlay = Image.new('RGBA', img_draw.size, (0, 0, 0, 0))
    draw_overlay = ImageDraw.Draw(overlay)
    
    try:
        font = ImageFont.load_default()
    except IOError:
        font = ImageFont.load_default()
        
    # Extract text and coordinates using regex
    pattern = r'([^()]+)(\(\d+,\d+\),\(\d+,\d+\))'
    matches = re.finditer(pattern, response)
    
    for match in matches:
        try:
            text = match.group(1).strip()
            coords = match.group(2)
            
            # Parse coordinates
            coord_pattern = r'\((\d+),(\d+)\)'
            coord_matches = re.findall(coord_pattern, coords)
            
            if len(coord_matches) == 2:
                x1, y1 = int(coord_matches[0][0]), int(coord_matches[0][1])
                x2, y2 = int(coord_matches[1][0]), int(coord_matches[1][1])
                
                # Generate random color
                color = (np.random.randint(0, 200), 
                        np.random.randint(0, 200),
                        np.random.randint(0, 255))
                color_alpha = color + (20,)
                
                # Draw rectangle and overlay
                draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
                draw_overlay.rectangle([x1, y1, x2, y2], 
                                    fill=color_alpha,
                                    outline=(0, 0, 0, 0))
                
                # Draw text label
                text_x = x1
                text_y = max(0, y1 - 15)
                text_bbox = draw.textbbox((0, 0), text, font=font)
                text_width = text_bbox[2] - text_bbox[0]
                text_height = text_bbox[3] - text_bbox[1]
                
                draw.rectangle([text_x, text_y, 
                              text_x + text_width, text_y + text_height],
                             fill=(255, 255, 255, 30))
                draw.text((text_x, text_y), text, font=font, fill=color)
                
        except Exception as e:
            print(f"Error drawing box: {str(e)}")
            continue
            
    # Combine image with overlay
    img_draw.paste(overlay, (0, 0), overlay)
    return img_draw

def main():
    """Main function to process images and visualize results"""
    # Read JSONL file
    jsonl_path = "/path/to/test_data.jsonl"
    output_dir = "output_visualizations"
    image_root = "/path/to/image_root"
    
    os.makedirs(output_dir, exist_ok=True)
    
    # Read all lines from JSONL
    items = []
    with open(jsonl_path, 'r') as f:
        for line in f:
            items.append(json.loads(line.strip()))
            
    # Randomly select one item
    item = random.choice(items)
    
    # Get image path and response
    image_path = os.path.join(image_root, item["image_name"])
    response = clean_repeated_substrings(item["vllm-infer"])
    
    print(f"Processing image: {item['image_name']}")
    
    # Load and process image
    image = Image.open(image_path)
    image_width, image_height = image.size
    
    # Process response coordinates
    processed_response = process_spotting_response(response, image_width, image_height)
    print("Original response:", response)
    print("Processed response:", processed_response)
    
    # Draw detection boxes
    result_image = draw_text_detection_boxes(image, processed_response)
    
    # Save result using original image name
    output_path = os.path.join(output_dir, item["image_name"])
    result_image.save(output_path)
    print(f"Visualization saved to {output_path}")

if __name__ == "__main__":
    main()