fix(predictions): TABLE, CELL & KEY_VALUE_SET blocks are not properly processed

AmplifyPlugins/Predictions/AWSPredictionsPlugin/Support/Utils/IdentifyResultTransformers.swift

@@ -67,6 +67,96 @@ class IdentifyResultTransformers {
 
     }
 
+    static func processChild(ids: [String],
+                             blockMap: [String: AWSTextractBlock]) -> String {
+        var keyText = ""
+        for keyId in ids {
+            let keyBlock = blockMap[keyId]
+            guard let keyBlockType = keyBlock?.blockType else {
+                continue
+            }
+            switch keyBlockType {
+            case .word:
+                if let text = keyBlock?.text {
+                    keyText += text + " "
+                }
+            default: break
+            }
+        }
+        return keyText
+    }
+
+    static func processValue(ids: [String],
+                             blockMap: [String: AWSTextractBlock]) -> (String, Bool) {
+        var valueText = ""
+        var valueSelected = false
+
+        // VALUE block has a CHILD list of IDs for the WORD block
+        for valueId in ids {
+            let valueBlock = blockMap[valueId]
+            guard let valueRelatioins = valueBlock?.relationships else {
+                continue
+            }
+            for valueRelation in valueRelatioins {
+                guard let ids = valueRelation.ids else {
+                    break
+                }
+                for id in ids {
+                    let wordBlock = blockMap[id]
+                    guard let wordValueBlockType = wordBlock?.blockType else {
+                        continue
+                    }
+                    switch wordValueBlockType {
+                    case .word:
+                        if let text = wordBlock?.text {
+                            valueText += text + " "
+                        }
+                    case .selectionElement:
+                        valueSelected = wordBlock?.selectionStatus == .selected ? true : false
+                    default: break
+                    }
+                }
+            }
+
+        }
+        return (valueText, valueSelected)
+    }
+
+    static func parseLineBlock(block: AWSTextractBlock) -> IdentifiedLine? {
+        guard let text = block.text,
+            let boundingBox = processBoundingBox(block.geometry?.boundingBox),
+            let polygon = processPolygon(block.geometry?.polygon) else {
+                return nil
+        }
+
+        return IdentifiedLine(text: text,
+                              boundingBox: boundingBox,
+                              polygon: polygon,
+                              page: Int(truncating: block.page ?? 0))
+    }
+
+    static func parseWordBlock(block: AWSTextractBlock) -> IdentifiedWord? {
+        guard let text = block.text,
+            let boundingBox = processBoundingBox(block.geometry?.boundingBox),
+            let polygon = processPolygon(block.geometry?.polygon) else {
+                return nil
+        }
+
+         return IdentifiedWord(text: text,
+                               boundingBox: boundingBox,
+                               polygon: polygon,
+                               page: Int(truncating: block.page ?? 0))
+    }
+
+    static func parseSelectionElementBlock(block: AWSTextractBlock) -> Selection? {
+        guard let boundingBox = processBoundingBox(block.geometry?.boundingBox),
+            let polygon = processPolygon(block.geometry?.polygon) else {
+                return nil
+        }
+        let selectionStatus = block.selectionStatus == .selected ? true : false
+        return Selection(boundingBox: boundingBox, polygon: polygon, isSelected: selectionStatus)
+    }
+
     // swiftlint:disable cyclomatic_complexity
     static func processLandmarks(_ rekognitionLandmarks: [AWSRekognitionLandmark]?) -> [Landmark] {
         var landmarks = [Landmark]()

AmplifyPlugins/Predictions/AWSPredictionsPlugin/Support/Utils/IdentifyTextResultTransformers.swift

@@ -53,79 +53,52 @@ class IdentifyTextResultTransformers: IdentifyResultTransformers {
     }
 
     static func processText(_ textractTextBlocks: [AWSTextractBlock]) -> IdentifyDocumentTextResult {
-
+        var fullText = ""
         var words = [IdentifiedWord]()
         var lines = [String]()
         var linesDetailed = [IdentifiedLine]()
         var selections = [Selection]()
-        var fullText = ""
         var tables = [Table]()
         var keyValues = [BoundedKeyValue]()
-        var blockMap = [String: AWSTextractBlock]()
         var tableBlocks = [AWSTextractBlock]()
         var keyValueBlocks = [AWSTextractBlock]()
+        var blockMap = [String: AWSTextractBlock]()
 
         for block in textractTextBlocks {
-            guard let text = block.text, let identifier = block.identifier else {
+            guard let identifier = block.identifier else {
                 continue
             }
 
-            guard let boundingBox = processBoundingBox(block.geometry?.boundingBox) else {
-                continue
-            }
-
-            guard let polygon = processPolygon(block.geometry?.polygon) else {
-                continue
-            }
-            let word = IdentifiedWord(text: text,
-                                      boundingBox: boundingBox,
-                                      polygon: polygon,
-                                      page: Int(truncating: block.page ?? 0))
-
-            let line = IdentifiedLine(text: text,
-                                      boundingBox: boundingBox,
-                                      polygon: polygon,
-                                      page: Int(truncating: block.page ?? 0))
+            blockMap[identifier] = block
 
             switch block.blockType {
             case .line:
-                lines.append(text)
-                linesDetailed.append(line)
+                if let line = parseLineBlock(block: block) {
+                    lines.append(line.text)
+                    linesDetailed.append(line)
+                }
             case .word:
-                fullText += text + " "
-                words.append(word)
-                blockMap[identifier] = block
+                if let word = parseWordBlock(block: block) {
+                    fullText += word.text + " "
+                    words.append(word)
+                    blockMap[identifier] = block
+                }
             case .selectionElement:
-                let selectionStatus = block.selectionStatus == .selected ? true : false
-                let selection = Selection(boundingBox: boundingBox, polygon: polygon, isSelected: selectionStatus)
-                selections.append(selection)
-                blockMap[identifier] = block
+                if let selection = parseSelectionElementBlock(block: block) {
+                    selections.append(selection)
+                    blockMap[identifier] = block
+                }
             case .table:
                 tableBlocks.append(block)
             case .keyValueSet:
                 keyValueBlocks.append(block)
-                blockMap[identifier] = block
             default:
-                blockMap[identifier] = block
-
-            }
-        }
-
-        if !tableBlocks.isEmpty {
-            for tableBlock in tableBlocks {
-                if let table = processTable(tableBlock, blockMap: blockMap) {
-                    tables.append(table)
-                }
+                continue
             }
         }
 
-        if !keyValueBlocks.isEmpty {
-            for keyValueBlock in keyValueBlocks where keyValueBlock.entityTypes?.contains("KEY") ?? false {
-                if let keyValue = processKeyValue(keyValueBlock, blockMap: blockMap) {
-                    keyValues.append(keyValue)
-                }
-            }
-        }
+        tables = processTables(tableBlocks: tableBlocks, blockMap: blockMap)
+        keyValues = processKeyValues(keyValueBlocks: keyValueBlocks, blockMap: blockMap)
 
         return IdentifyDocumentTextResult(
             fullText: fullText,
@@ -137,15 +110,47 @@ class IdentifyTextResultTransformers: IdentifyResultTransformers {
             keyValues: keyValues)
     }
 
-    static func processTable(_ block: AWSTextractBlock,
+    static func processTables(tableBlocks: [AWSTextractBlock],
+                              blockMap: [String: AWSTextractBlock]) -> [Table] {
+        var tables = [Table]()
+        for tableBlock in tableBlocks {
+            if let table = processTable(tableBlock, blockMap: blockMap) {
+                tables.append(table)
+            }
+        }
+        return tables
+    }
+
+    static func processKeyValues(keyValueBlocks: [AWSTextractBlock],
+                                 blockMap: [String: AWSTextractBlock]) -> [BoundedKeyValue] {
+        var keyValues =  [BoundedKeyValue]()
+        for keyValueBlock in keyValueBlocks where keyValueBlock.entityTypes?.contains("KEY") ?? false {
+            if let keyValue = processKeyValue(keyValueBlock, blockMap: blockMap) {
+                keyValues.append(keyValue)
+            }
+        }
+        return keyValues
+    }
+
+    // https://docs.aws.amazon.com/textract/latest/dg/how-it-works-tables.html
+    /**
+    * Converts a given Amazon Textract block into Amplify-compatible
+    * table object.
+    * @param block Textract text block
+    * @param blockMap map of Textract blocks by their IDs
+    * @return Amplify Table instance
+    */
+    static func processTable(_ tableBlock: AWSTextractBlock,
                              blockMap: [String: AWSTextractBlock]) -> Table? {
 
-        guard let relationships = block.relationships else {
+        guard let relationships = tableBlock.relationships else {
             return nil
         }
         var table = Table()
         var rows = Set<Int>()
         var cols = Set<Int>()
+
+        // Each TABLE block contains CELL blocks
         for tableRelation in relationships {
             guard let ids = tableRelation.ids else {
                 continue
@@ -162,13 +167,13 @@ class IdentifyTextResultTransformers: IdentifyResultTransformers {
                 let row = Int(truncating: rowIndex) - 1
                 let col = Int(truncating: colIndex) - 1
                 if !rows.contains(row) {
-                rows.insert(row)
+                    rows.insert(row)
                 }
                 if !cols.contains(col) {
-                cols.insert(col)
+                    cols.insert(col)
                 }
-                if let cell = constructTableCell(cellBlock) {
-                table.cells.append(cell)
+                if let cell = constructTableCell(cellBlock, blockMap) {
+                    table.cells.append(cell)
                 }
             }
 
@@ -178,27 +183,42 @@ class IdentifyTextResultTransformers: IdentifyResultTransformers {
         return table
     }
 
-    static func constructTableCell(_ block: AWSTextractBlock?) -> Table.Cell? {
-        guard let blockType = block?.blockType,
-            let selectionStatus = block?.selectionStatus,
-            let text = block?.text,
+    static func constructTableCell(_ block: AWSTextractBlock?, _ blockMap: [String: AWSTextractBlock]) -> Table.Cell? {
+        guard let selectionStatus = block?.selectionStatus,
             let rowSpan = block?.rowSpan,
             let columnSpan = block?.columnSpan,
             let geometry = block?.geometry,
             let textractBoundingBox = geometry.boundingBox,
-            let texttractPolygon = geometry.polygon
+            let texttractPolygon = geometry.polygon,
+            let relationships = block?.relationships
             else {
                 return nil
         }
+
         var words = ""
         var isSelected = false
 
-        switch blockType {
-        case .word:
-             words += text + " "
-        case .selectionElement:
-            isSelected = selectionStatus == .selected ? true : false
-        default: break
+        // Each CELL block consists of WORD and/or SELECTION_ELEMENT blocks
+        for cellRelation in relationships {
+            guard let ids = cellRelation.ids else {
+                continue
+            }
+
+            for wordId in ids {
+                let wordBlock = blockMap[wordId]
+
+                switch wordBlock?.blockType {
+                case .word:
+                    guard let text = wordBlock?.text else {
+                        return nil
+                    }
+                    words += text + " "
+                case .selectionElement:
+                    isSelected = selectionStatus == .selected ? true : false
+                default:
+                    break
+                }
+            }
         }
 
         guard let boundingBox = processBoundingBox(textractBoundingBox) else {
@@ -218,6 +238,14 @@ class IdentifyTextResultTransformers: IdentifyResultTransformers {
         return cell
     }
 
+    // https://docs.aws.amazon.com/textract/latest/dg/how-it-works-kvp.html
+    /**
+    * Converts a given Amazon Textract block into Amplify-compatible
+    * key-value pair feature. Returns null if not a valid table.
+    * @param block Textract text block
+    * @param blockMap map of Textract blocks by their IDs
+    * @return Amplify KeyValue instance
+    */
     static func processKeyValue(_ keyBlock: AWSTextractBlock?,
                                 blockMap: [String: AWSTextractBlock]) -> BoundedKeyValue? {
         var keyText = ""
@@ -229,32 +257,20 @@ class IdentifyTextResultTransformers: IdentifyResultTransformers {
             return nil
         }
 
+        // KEY_VALUE_SET block contains CHILD and VALUE entity type blocks
         for keyBlockRelationship in relationships {
-            guard let text = keyBlock.text,
-                let ids = keyBlockRelationship.ids else {
-                    continue
+
+            guard let ids = keyBlockRelationship.ids else {
+                continue
             }
+
             switch keyBlockRelationship.types {
-            case .child where keyBlock.blockType == .word:
-                     keyText += text + " "
+            case .child:
+                keyText = processChild(ids: ids, blockMap: blockMap)
             case .value:
-                for valueId in ids {
-                     let valueBlock = blockMap[valueId]
-                    guard let valueBlockType = valueBlock?.blockType else {
-                        continue
-                    }
-                    switch valueBlockType {
-                    case .word:
-
-                        if let text = valueBlock?.text {
-                        valueText += text + " "
-                        }
-                    case .selectionElement:
-                        valueSelected = keyBlock.selectionStatus == .selected ? true : false
-                    default: break
-                    }
-                }
-
+                let valueResult = processValue(ids: ids, blockMap: blockMap)
+                valueText = valueResult.0
+                valueSelected = valueResult.1
             default:
                 break
             }
@@ -269,9 +285,9 @@ class IdentifyTextResultTransformers: IdentifyResultTransformers {
         }
 
         return BoundedKeyValue(key: keyText,
-                        value: valueText,
-                        isSelected: valueSelected,
-                        boundingBox: boundingBox,
-                        polygon: polygon)
+                               value: valueText,
+                               isSelected: valueSelected,
+                               boundingBox: boundingBox,
+                               polygon: polygon)
     }
 }