Update app.py

app.py

@@ -20,10 +20,6 @@ from jinja2 import Template
 import re
 import time
 
-# Add parent directory to path to import seg_rec_kraken
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from seg_rec_kraken import seg_rec_image, PAGEXML_TEMPLATE
-
 # Configure logging
 logging.basicConfig(level=logging.WARNING)
 logging.getLogger('kraken').setLevel(logging.WARNING)
@@ -36,11 +32,127 @@ MODELS_DIR = Path("models")
 SEG_MODELS_DIR = MODELS_DIR / "seg"
 REC_MODELS_DIR = MODELS_DIR / "rec"
 
+# Embedded template
+PAGEXML_TEMPLATE = '''{%+ macro render_line(line) +%}
+            <TextLine id="{{ line.id }}" {% if line.tags and "type" in line.tags %}custom="structure {type:{{ line.tags["type"] }};}"{% endif %}>
+                {% if line.boundary %}
+                <Coords points="{% for point in line.boundary %}{{ point|join(',') }}{% if not loop.last %} {% endif %}{% endfor %}"/>
+                {% endif %}
+                {% if line.baseline %}
+                <Baseline points="{% for point in line.baseline %}{{ point|join(',') }}{% if not loop.last %} {% endif %}{% endfor %}"/>
+                {% endif %}
+                {% if line.text is string %}
+                    <TextEquiv{% if line.confidences|length %} conf="{{ (line.confidences|sum / line.confidences|length)|round(4) }}"{% endif %}><Unicode>{{ line.text|e }}</Unicode></TextEquiv>
+                {% else %}
+                {% for segment in line.recognition %}
+                <Word id="segment_{{ segment.index }}">
+                    {% if segment.boundary %}
+                    <Coords points="{% for point in segment.boundary %}{{ point|join(',') }}{% if not loop.last %} {% endif %}{% endfor %}"/>
+                    {% else %}
+                    <Coords points="{{ segment.bbox[0] }},{{ segment.bbox[1] }} {{ segment.bbox[0] }},{{ segment.bbox[3] }} {{ segment.bbox[2] }},{{ segment.bbox[3] }} {{ segment.bbox[2] }},{{ segment.bbox[1] }}"/>
+                    {% endif %}
+                {% for char in segment.recognition %}
+                    <Glyph id="char_{{ char.index }}">
+                        <Coords points="{% for point in char.boundary %}{{ point|join(',') }}{% if not loop.last %} {% endif %}{% endfor %}"/>
+                        <TextEquiv conf="{{ char.confidence|round(4) }}"><Unicode>{{ char.text|e }}</Unicode></TextEquiv>
+                    </Glyph>
+                {% endfor %}
+                    <TextEquiv conf="{{ (segment.confidences|sum / segment.confidences|length)|round(4) }}"><Unicode>{{ segment.text|e }}</Unicode></TextEquiv>
+                </Word>
+                {% endfor %}
+                {%+ if line.confidences|length %}<TextEquiv conf="{{ (line.confidences|sum / line.confidences|length)|round(4) }}"><Unicode>{% for segment in line.recognition %}{{ segment.text|e }}{% endfor %}</Unicode></TextEquiv>{% endif +%}
+                {% endif %}
+            </TextLine>
+{%+ endmacro %}
+<?xml version="1.0" encoding="UTF-8"?>
+<PcGts xmlns="http://schema.primaresearch.org/PAGE/gts/pagecontent/2019-07-15" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://schema.primaresearch.org/PAGE/gts/pagecontent/2019-07-15 http://schema.primaresearch.org/PAGE/gts/pagecontent/2019-07-15/pagecontent.xsd">
+    <Metadata>
+        <Creator>kraken {{ metadata.version }}</Creator>
+        <Created>{{ page.date }}</Created>
+        <LastChange>{{ page.date }}</LastChange>
+    </Metadata>
+    <Page imageFilename="{{ page.name }}" imageWidth="{{ page.size[0] }}" imageHeight="{{ page.size[1] }}" {% if page.base_dir %}readingDirection="{{ page.base_dir }}"{% endif %}>
+    {% for entity in page.entities %}
+        {% if entity.type == "region" %}
+        {% if loop.previtem and loop.previtem.type == 'line' %}
+        </TextRegion>
+        {% endif %}
+        <TextRegion id="{{ entity.id }}" {% if entity.tags and "type" in entity.tags %}custom="structure {type:{{ entity.tags["type"] }};}"{% endif %}>
+            {% if entity.boundary %}<Coords points="{% for point in entity.boundary %}{{ point|join(',') }}{% if not loop.last %} {% endif %}{% endfor %}"/>{% endif %}
+            {%- for line in entity.lines -%}
+            {{ render_line(line) }}
+            {%- endfor %}
+        </TextRegion>
+        {% else %}
+        {% if not loop.previtem or loop.previtem.type != 'line' %}
+        <TextRegion id="textblock_{{ loop.index }}">
+            <Coords points="0,0 0,{{ page.size[1] }} {{ page.size[0] }},{{ page.size[1] }} {{ page.size[0] }},0"/>
+        {% endif %}
+            {{ render_line(entity) }}
+        {% if loop.last %}
+        </TextRegion>
+        {% endif %}
+        {% endif %}
+    {% endfor %}
+    </Page>
+</PcGts>'''
+
 # Create Jinja environment
 TEMPLATE_DIR = Path("templates")
 TEMPLATE_DIR.mkdir(exist_ok=True)
 _ENV = Environment(loader=FileSystemLoader(str(TEMPLATE_DIR)))
 
+def seg_rec_image(image_path, seg_model, rec_model, output_dir=None):
+    try:
+        im = Image.open(image_path)
+        baseline_seg = blla.segment(im, model=seg_model)
+
+        # Run recognition and collect full BaselineOCRRecord objects
+        pred_it = rpred.rpred(network=rec_model, im=im, bounds=baseline_seg, pad=16)
+        records = [record for record in pred_it]
+
+        # Attach recognition results to segmentation lines
+        for line, rec_line in zip(baseline_seg.lines, records):
+            # Debug logging for recognition results
+            logger.debug(f'Recognition result - Prediction: {rec_line.prediction}')
+            logger.debug(f'Recognition result - Confidences: {rec_line.confidences}')
+
+            # Ensure the line has both prediction and confidence values
+            line.prediction = rec_line.prediction
+            line.text = rec_line.prediction  # Set text field for serialization
+
+            # Store per-character confidences
+            line.confidences = rec_line.confidences  # Keep the list of confidences
+
+            # Debug logging for line object
+            logger.debug(f'Line {line.id} - Prediction: {line.prediction}')
+            logger.debug(f'Line {line.id} - Confidences: {line.confidences}')
+
+        # Construct PAGE-XML segmentation only data
+        pagexml_seg_only = serialization.serialize(baseline_seg, image_size=im.size, template='pagexml', sub_line_segmentation=False)
+
+        # Serialize with recognition results
+        pagexml = serialization.serialize(baseline_seg,
+                                        image_size=im.size,
+                                        template='custom_pagexml',
+                                        template_source='custom',
+                                        sub_line_segmentation=False)
+
+        base_name = os.path.splitext(os.path.basename(image_path))[0]
+        if output_dir:
+            os.makedirs(output_dir, exist_ok=True)
+            output_path = os.path.join(output_dir, base_name + '.xml')
+        else:
+            output_path = os.path.splitext(image_path)[0] + '.xml'
+
+        with open(output_path, 'w') as fp:
+            fp.write(pagexml)
+        print(f"✅ Segmented/recognized: {os.path.basename(image_path)} → {os.path.basename(output_path)}")
+    except Exception as e:
+        print(f"❌ Failed to process {image_path}: {e}")
+        import traceback
+        traceback.print_exc()
+
 # Create template files
 def create_templates():
     """Create Jinja templates for visualization."""