test_skew_correction.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. import unittest
  2. import numpy as np
  3. from PIL import Image
  4. import sys
  5. import os
  6. sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
  7. from ocr_verify.table_line_generator.table_line_generator import TableLineGenerator
  8. class TestSkewCorrection(unittest.TestCase):
  9. def test_estimate_skew_from_rows(self):
  10. # Mock data: 3 rows, each with 3 cells, skewed by ~5 degrees
  11. # tan(5 deg) approx 0.087
  12. skew_slope = 0.087
  13. text_boxes = []
  14. for row in range(3):
  15. y_base = row * 100
  16. for col in range(3):
  17. x_center = col * 100 + 50
  18. y_center = y_base + x_center * skew_slope + 50
  19. # Create a 40x20 box around center
  20. bbox = [
  21. x_center - 20, y_center - 10,
  22. x_center + 20, y_center + 10
  23. ]
  24. text_boxes.append({
  25. 'row': row + 1,
  26. 'col': col + 1,
  27. 'bbox': bbox,
  28. 'text': f"R{row}C{col}"
  29. })
  30. ocr_data = {'actual_rows': 3, 'text_boxes': text_boxes}
  31. # Create dummy image
  32. img = Image.new('RGB', (400, 400), color='white')
  33. generator = TableLineGenerator(img, ocr_data, auto_correct_skew=False)
  34. estimated_angle = generator._estimate_skew_from_rows(text_boxes)
  35. print(f"Estimated Angle: {estimated_angle:.4f} degrees")
  36. self.assertAlmostEqual(estimated_angle, 5.0, delta=0.5)
  37. if __name__ == '__main__':
  38. unittest.main()