import unittest import numpy as np from PIL import Image import sys import os sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))) from ocr_verify.table_line_generator.table_line_generator import TableLineGenerator class TestSkewCorrection(unittest.TestCase): def test_estimate_skew_from_rows(self): # Mock data: 3 rows, each with 3 cells, skewed by ~5 degrees # tan(5 deg) approx 0.087 skew_slope = 0.087 text_boxes = [] for row in range(3): y_base = row * 100 for col in range(3): x_center = col * 100 + 50 y_center = y_base + x_center * skew_slope + 50 # Create a 40x20 box around center bbox = [ x_center - 20, y_center - 10, x_center + 20, y_center + 10 ] text_boxes.append({ 'row': row + 1, 'col': col + 1, 'bbox': bbox, 'text': f"R{row}C{col}" }) ocr_data = {'actual_rows': 3, 'text_boxes': text_boxes} # Create dummy image img = Image.new('RGB', (400, 400), color='white') generator = TableLineGenerator(img, ocr_data, auto_correct_skew=False) estimated_angle = generator._estimate_skew_from_rows(text_boxes) print(f"Estimated Angle: {estimated_angle:.4f} degrees") self.assertAlmostEqual(estimated_angle, 5.0, delta=0.5) if __name__ == '__main__': unittest.main()