Coverage for collectra/tasks/ocr.py: 75.32%

1__all__ = ["OCRSuraya", "LineDetectorSurya"] 

2 

3from dotenv import load_dotenv 

4 

5from collectra.tasks.base import Task 

6from collectra.types.images import Image, ImageCrop 

7from collectra.types.texts import Text 

8 

9load_dotenv() 

10 

11 

12class OCRSuraya(Task): 

13 def __init__(self, name: str, **kwargs) -> None: 

14 """Initialize the task without loading Surya model weights.""" 

15 super().__init__(name, **kwargs) 

16 self.detection_predictor = None 

17 self.recognition_predictor = None 

18 

19 def _init_predictors(self) -> None: 

20 """Load Surya predictors on first use.""" 

21 if self.detection_predictor is not None and self.recognition_predictor is not None: 

22 return 

23 from surya.detection import DetectionPredictor 

24 from surya.foundation import FoundationPredictor 

25 from surya.recognition import RecognitionPredictor 

26 

27 self.detection_predictor = DetectionPredictor() 

28 foundation_predictor = FoundationPredictor() 

29 self.recognition_predictor = RecognitionPredictor(foundation_predictor) 

30 

31 def input_type(self) -> type | tuple: 

32 """Define the expected input types for this LLM task. 

33 

34 Returns: 

35 type | tuple: Tuple of str and Path types for text and file inputs. 

36 """ 

37 return Image 

38 

39 def output_type(self) -> type | tuple: 

40 """Define the expected output types for this LLM task. 

41 

42 Returns: 

43 type | tuple: String type for generated text outputs. 

44 """ 

45 return Text 

46 

47 def run(self, *images: Image) -> list[Text]: 

48 """Recognize text in each image and return Text artefacts.""" 

49 self._init_predictors() 

50 results: list[Text] = [] 

51 for image in images: 

52 if not isinstance(image, Image): 

53 raise TypeError(f"image must be Image, got {type(image)}") 

54 predictions = self.recognition_predictor( 

55 [image.pil()], det_predictor=self.detection_predictor 

56 ) 

57 text = "\n".join(line.text for line in predictions[0].text_lines) 

58 results.append(Text(name=self.get_output_name(), data=text)) 

59 return results 

60 

61 

62class LineDetectorSurya(Task): 

63 def __init__( 

64 self, 

65 name: str, 

66 merge_horizontal: bool = False, 

67 min_height: int = 0, 

68 **kwargs, 

69 ) -> None: 

70 """Initialize the task without loading Surya model weights.""" 

71 super().__init__( 

72 name, 

73 merge_horizontal=merge_horizontal, 

74 min_height=min_height or 0, 

75 **kwargs, 

76 ) 

77 self.detection_predictor = None 

78 

79 def _init_predictor(self) -> None: 

80 """Load the Surya detection predictor on first use.""" 

81 if self.detection_predictor is not None: 

82 return 

83 from surya.detection import DetectionPredictor 

84 

85 self.detection_predictor = DetectionPredictor() 

86 

87 def input_type(self) -> type | tuple: 

88 """Define the expected input types for this LLM task. 

89 

90 Returns: 

91 type | tuple: Tuple of str and Path types for text and file inputs. 

92 """ 

93 return Image 

94 

95 def output_type(self) -> type | tuple: 

96 """Define the expected output types for this LLM task. 

97 

98 Returns: 

99 type | tuple: String type for generated text outputs. 

100 """ 

101 return ImageCrop 

102 

103 def run(self, *images: Image) -> list[ImageCrop]: 

104 """Detect text lines and return their image crops.""" 

105 self._init_predictor() 

106 results: list[ImageCrop] = [] 

107 for image in images: 

108 if not isinstance(image, Image): 

109 raise TypeError(f"image must be Image, got {type(image)}") 

110 pil_image = image.pil() 

111 predictions = self.detection_predictor([pil_image]) 

112 

113 # Sort bounding boxes vertically 

114 bounding_boxes = sorted( 

115 [polygon_box.bbox for polygon_box in predictions[0].bboxes], 

116 key=lambda bbox: bbox[1], 

117 ) 

118 

119 if len(bounding_boxes) == 0 or ( 

120 self.min_height and image.height <= self.min_height 

121 ): 

122 crop = image.make_crop(0.5, 0.5, 1.0, 1.0) 

123 results.append(crop) 

124 elif self.merge_horizontal: 

125 bbox = bounding_boxes[0] 

126 top = bbox[1] 

127 bottom = bbox[3] 

128 

129 for bbox in bounding_boxes: 

130 if bbox[1] >= bottom: 

131 crop = image.make_crop_bounding_box( 

132 0, top, image.width, bottom, min_height=self.min_height 

133 ) 

134 results.append(crop) 

135 top = bbox[1] 

136 bottom = bbox[3] 

137 else: 

138 top = min(top, bbox[1]) 

139 bottom = max(bottom, bbox[3]) 

140 

141 crop = image.make_crop_bounding_box(0, top, image.width, bottom) 

142 results.append(crop) 

143 else: 

144 for bbox in bounding_boxes: 

145 crop = image.make_crop_bounding_box( 

146 *bbox, min_height=self.min_height 

147 ) 

148 results.append(crop) 

149 

150 return results