Merge remote-tracking branch 'bertsky/extract-regions-coco'

Author: Robert Sachunsky

ocrd_segment/extract_pages.py

@@ -127,6 +127,7 @@ def process(self):
         Create two JSON files with region types and coordinates: one (page-wise) in
         our custom format and one (global) in MS-COCO.
         
+        \b
         The output file group may be given as a comma-separated list to separate
         these 3 kinds of images. If fewer than 3 fileGrps are specified, they will
         share the same fileGrp (and directory). In particular, write files as follows:
@@ -147,8 +148,9 @@ def process(self):
             each layer and segment is superimposed (alpha blended) onto the previous one,
             starting with the above raw image.
         
+        \b
         In addition, write a file for all pages at once:
-        * in the third (or first) output file group (directory):
+        * in the third (or second or only) output file group (directory):
           - output_file_grp + '.coco.json': region coordinates/classes (MS-COCO format)
           - output_file_grp + '.colordict.json': the used ``colordict``
         

ocrd_segment/extract_regions.py

@@ -1,20 +1,23 @@
 from __future__ import absolute_import
 
+import os
 import json
+import numpy as np
 
 from ocrd_utils import (
     getLogger,
     make_file_id,
     assert_file_grp_cardinality,
     coordinates_of_segment,
+    xywh_from_polygon,
     polygon_from_points,
     MIME_TO_EXT
 )
 from ocrd_modelfactory import page_from_file
 from ocrd import Processor
 
 from .config import OCRD_TOOL
-from .extract_pages import CLASSES
+from .extract_pages import CLASSES, segment_poly
 
 TOOL = 'ocrd-segment-extract-regions'
 
@@ -40,7 +43,10 @@ def process(self):
         If ``transparency`` is true, then also add an alpha channel which is
         fully transparent outside of the mask.
         
-        Create a JSON file with:
+        Create two JSON files with region types and coordinates: one (page-wise)
+        in our custom format and one (global) in MS-COCO.
+        
+        The custom JSON files contain:
         * the IDs of the region and its parents,
         * the region's coordinates relative to the region image,
         * the region's absolute coordinates,
@@ -62,6 +68,7 @@ def process(self):
         * ID + '.bin.png': region image (if the workflow provides binarized images)
         * ID + '.nrm.png': region image (if the workflow provides grayscale-normalized images)
         * ID + '.json': region metadata.
+        * output_file_grp + '.coco.json'
         """
         LOG = getLogger('processor.ExtractRegions')
         assert_file_grp_cardinality(self.input_file_grp, 1)
@@ -72,12 +79,32 @@ def process(self):
         if self.parameter["classes"]:
             selected_classes = self.parameter["classes"]
             classes = { region: classes[region] for region in selected_classes }
+        # COCO: init data structures
+        images = list()
+        annotations = list()
+        categories = list()
+        i = 0
+        for cat, color in classes.items():
+            # COCO format does not allow alpha channel
+            color = (int(color[0:2], 16),
+                     int(color[2:4], 16),
+                     int(color[4:6], 16))
+            try:
+                supercat, name = cat.split(':')
+            except ValueError:
+                name = cat
+                supercat = ''
+            categories.append(
+                {'id': i, 'name': name, 'supercategory': supercat,
+                 'source': 'PAGE', 'color': color})
+            i += 1
+        i = 0 # subregion count (i.e. annotation id)
+        j = 0 # region count (i.e. image id)
         # pylint: disable=attribute-defined-outside-init
         for n, input_file in enumerate(self.input_files):
             page_id = input_file.pageId or input_file.ID
             LOG.info("INPUT FILE %i / %s", n, page_id)
             pcgts = page_from_file(self.workspace.download_file(input_file))
-            self.add_metadata(pcgts)
             page = pcgts.get_Page()
             page_image, page_coords, page_image_info = self.workspace.image_from_page(
                 page, page_id,
@@ -92,9 +119,9 @@ def process(self):
             ptype = page.get_type()
 
             regions = dict()
-            for name in classes.keys():
+            for name in classes:
                 if not name or not name.endswith("Region"):
-                    # no subtypes here
+                    # only top-level regions here
                     continue
                 regions[name] = getattr(page, 'get_' + name)()
             for rtype, rlist in regions.items():
@@ -110,6 +137,7 @@ def process(self):
                         subrtype = region.get_type()
                     else:
                         subrtype = None
+                    j += 1
                     description['subtype'] = subrtype
                     description['coords_rel'] = coordinates_of_segment(
                         region, region_image, region_coords).tolist()
@@ -164,7 +192,50 @@ def process(self):
                         extension = '.nrm'
                     else:
                         extension = '.raw'
-
+                    subregions = dict()
+                    for name in classes:
+                        if not name or ':' in name:
+                            # no subtypes here
+                            continue
+                        if not hasattr(region, 'get_' + name):
+                            continue
+                        subregions[name] = getattr(region, 'get_' + name)()
+                    for subrtype, subrlist in subregions.items():
+                        for subregion in subrlist:
+                            poly = segment_poly(page_id, subregion, region_coords)
+                            if not poly:
+                                continue
+                            polygon = np.array(poly.exterior.coords, np.int)[:-1].tolist()
+                            xywh = xywh_from_polygon(polygon)
+                            area = poly.area
+                            if subrtype in ['TextRegion', 'ChartRegion', 'GraphicRegion']:
+                                subsubrtype = subregion.get_type()
+                            else:
+                                subsubrtype = None
+                            if subsubrtype:
+                                subrtype0 = subrtype + ':' + subsubrtype
+                            else:
+                                subrtype0 = subrtype
+                            description.setdefault('regions', []).append(
+                                { 'type': subrtype,
+                                  'subtype': subsubrtype,
+                                  'coords': polygon,
+                                  'area': area,
+                                  'region.ID': subregion.id
+                                })
+                            # COCO: add annotations
+                            i += 1
+                            annotations.append(
+                                {'id': i, 'image_id': j,
+                                 'category_id': next((cat['id'] for cat in categories if cat['name'] == subsubrtype),
+                                                     next((cat['id'] for cat in categories if cat['name'] == subrtype))),
+                                 'segmentation': np.array(poly.exterior.coords, np.int)[:-1].reshape(1, -1).tolist(),
+                                 'area': area,
+                                 'bbox': [xywh['x'], xywh['y'], xywh['w'], xywh['h']],
+                                 'iscrowd': 0})
+                            
+                                  
+                    
                     file_id = make_file_id(input_file, self.output_file_grp) + '_' + region.id + extension
                     file_path = self.workspace.save_image_file(
                         region_image,
@@ -179,3 +250,25 @@ def process(self):
                         pageId=input_file.pageId,
                         mimetype='application/json',
                         content=json.dumps(description))
+                    # COCO: add image
+                    images.append({
+                        'id': j,
+                        # all exported coordinates are relative to the cropped region:
+                        # -> use that for reference
+                        'file_name': file_path,
+                        # -> use its size
+                        'width': region_image.width,
+                        'height': region_image.height})
+        # COCO: write result
+        file_id = self.output_file_grp + '.coco.json'
+        LOG.info('Writing COCO result file "%s"', file_id)
+        self.workspace.add_file(
+            ID=file_id,
+            file_grp=self.output_file_grp,
+            local_filename=os.path.join(self.output_file_grp, file_id),
+            mimetype='application/json',
+            pageId=None,
+            content=json.dumps(
+                {'categories': categories,
+                 'images': images,
+                 'annotations': annotations}))