✨ support for custom API classification field

lib/mindee/documents/custom.rb

@@ -7,7 +7,7 @@ module Mindee
   # Custom document object.
   class CustomDocument < Document
     # All fields in the document
-    # @return [Hash<Symbol, Mindee::ListField>]
+    # @return [Hash<Symbol, Mindee::ListField | Mindee::ClassificationField>]
     attr_reader :fields
 
     # @param document_type [String]
@@ -19,7 +19,7 @@ def initialize(document_type, prediction, input_file: nil, page_id: nil)
       @fields = {}
       prediction.each do |field_name, field_prediction|
         field_sym = field_name.to_sym
-        complete_field = ListField.new(field_prediction, page_id)
+        complete_field = build_field(field_prediction, page_id)
 
         # Add the field to the `fields` array
         @fields[field_sym] = complete_field
@@ -40,5 +40,17 @@ def to_s
       out_str << "\n----------------------"
       out_str
     end
+
+    private
+
+    def build_field(field_prediction, page_id)
+      # Currently two types of fields possible in a custom API response:
+      # fields having a list of values, and classification fields.
+      # Here we use the fact that only value lists have the 'values' attribute.
+
+      return ListField.new(field_prediction, page_id) if field_prediction['values']
+
+      ClassificationField.new(field_prediction)
+    end
   end
 end

lib/mindee/fields.rb

@@ -8,4 +8,4 @@
 require_relative 'fields/orientation'
 require_relative 'fields/payment_details'
 require_relative 'fields/tax'
-require_relative 'fields/list_field'
+require_relative 'fields/custom_docs'

lib/mindee/fields/list_field.rb → lib/mindee/fields/custom_docs.rb

@@ -1,6 +1,19 @@
 # frozen_string_literal: true
 
 module Mindee
+  # Document classification (custom docs)
+  class ClassificationField
+    # @param prediction [Hash]
+    def initialize(prediction)
+      @value = prediction['value']
+      @confidence = prediction['confidence']
+    end
+
+    def to_s
+      @value
+    end
+  end
+
   # Field in a list.
   class ListFieldItem
     # The confidence score, value will be between 0.0 and 1.0

spec/data/custom/response/complete.json

@@ -571,6 +571,10 @@
         }
       ],
       "prediction": {
+        "doc_type": {
+          "confidence": 0.53,
+          "value": "type_b"
+        },
         "date_normal": {
           "confidence": 0.99,
           "page_id": 0,
@@ -907,7 +911,8 @@
           "date_us",
           "phone",
           "email",
-          "url"
+          "url",
+          "doc_type"
         ],
         "name": "ianare/field_test",
         "type": "constructed",

spec/data/custom/response/doc_to_string.txt

@@ -1,5 +1,6 @@
 ----- field_test -----
 Filename:
+doc_type: type_b
 date_normal: 2020-12-17
 date_us: 2020-12-17
 email: contact@mindee.com

spec/data/custom/response/empty.json

@@ -66,6 +66,10 @@
         }
       ],
       "prediction": {
+        "doc_type": {
+          "confidence": 0.0,
+          "value": null
+        },
         "date_normal": {
           "confidence": 0.0,
           "page_id": 0,

spec/documents.rb → spec/documents_spec.rb

@@ -135,7 +135,7 @@ def load_json(file_path)
       prediction = response['document']['inference']['prediction']
       document = Mindee::CustomDocument.new('field_test', prediction)
       expect(document.document_type).to eq('field_test')
-      expect(document.fields.length).to eq(10)
+      expect(document.fields.length).to eq(11)
     end
 
     it 'should load a complete document prediction' do