Run pipeline on one image. Returns output XML path or None.
(self,
image_path: str,
output_dir: str = None,
with_refinement: bool = False,
with_text: bool = True,
groups: List[PromptGroup] = None)
| 140 | return self._refinement_processor |
| 141 | |
| 142 | def process_image(self, |
| 143 | image_path: str, |
| 144 | output_dir: str = None, |
| 145 | with_refinement: bool = False, |
| 146 | with_text: bool = True, |
| 147 | groups: List[PromptGroup] = None) -> Optional[str]: |
| 148 | """Run pipeline on one image. Returns output XML path or None.""" |
| 149 | print(f"\n{'='*60}") |
| 150 | print(f"Processing: {image_path}") |
| 151 | print(f"{'='*60}") |
| 152 | |
| 153 | # Output directory |
| 154 | if output_dir is None: |
| 155 | output_dir = self.config.get('paths', {}).get('output_dir', './output') |
| 156 | |
| 157 | img_stem = Path(image_path).stem |
| 158 | img_output_dir = os.path.join(output_dir, img_stem) |
| 159 | os.makedirs(img_output_dir, exist_ok=True) |
| 160 | |
| 161 | print("\n[0] Preprocess...") |
| 162 | context = ProcessingContext( |
| 163 | image_path=image_path, |
| 164 | output_dir=img_output_dir |
| 165 | ) |
| 166 | context.intermediate_results['original_image_path'] = image_path |
| 167 | context.intermediate_results['was_upscaled'] = False |
| 168 | context.intermediate_results['upscale_factor'] = 1.0 |
| 169 | |
| 170 | try: |
| 171 | if with_text and self.text_restorer is not None: |
| 172 | print("\n[1] Text extraction (OCR)...") |
| 173 | try: |
| 174 | text_xml_content = self.text_restorer.process(image_path) |
| 175 | text_output_path = os.path.join(img_output_dir, "text_only.drawio") |
| 176 | with open(text_output_path, 'w', encoding='utf-8') as f: |
| 177 | f.write(text_xml_content) |
| 178 | context.intermediate_results['text_xml'] = text_xml_content |
| 179 | print(f" Saved: {text_output_path}") |
| 180 | except Exception as e: |
| 181 | print(f" Text step failed: {e}") |
| 182 | print(" Continuing without text...") |
| 183 | elif with_text: |
| 184 | print("\n[1] Text extraction (skipped - deps)") |
| 185 | else: |
| 186 | print("\n[1] Text extraction (skipped)") |
| 187 | |
| 188 | print("\n[2] Segmentation (SAM3)...") |
| 189 | |
| 190 | if groups: |
| 191 | # Extract by group |
| 192 | all_elements = [] |
| 193 | for group in groups: |
| 194 | result = self.sam3_extractor.extract_by_group(context, group) |
| 195 | all_elements.extend(result.elements) |
| 196 | for i, elem in enumerate(all_elements): |
| 197 | elem.id = i |
| 198 | context.elements = all_elements |
| 199 | context.canvas_width = result.canvas_width |
no test coverage detected