Skip to content

Commit f7f99fa

Browse files
authored
✨ add OCR output to CLI (#228)
1 parent f7a37c5 commit f7f99fa

File tree

1 file changed

+54
-20
lines changed

1 file changed

+54
-20
lines changed

src/main/java/com/mindee/CommandLineInterface.java

+54-20
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,20 @@
77
import com.mindee.input.PageOptions;
88
import com.mindee.input.PageOptionsOperation;
99
import com.mindee.parsing.common.AsyncPredictResponse;
10-
import com.mindee.parsing.common.Document;
1110
import com.mindee.parsing.common.Inference;
1211
import com.mindee.parsing.common.PredictResponse;
12+
import com.mindee.parsing.common.ocr.Ocr;
1313
import com.mindee.product.custom.CustomV1;
1414
import com.mindee.product.generated.GeneratedV1;
1515
import java.io.File;
1616
import java.io.IOException;
1717
import java.lang.reflect.Method;
1818
import java.util.ArrayList;
19+
import java.util.Arrays;
1920
import java.util.List;
2021
import java.util.Objects;
2122
import java.util.concurrent.Callable;
23+
import java.util.stream.Collectors;
2224
import picocli.CommandLine;
2325
import picocli.CommandLine.Command;
2426
import picocli.CommandLine.Model.CommandSpec;
@@ -66,7 +68,8 @@ private enum OutputChoices { summary, full, raw }
6668
description = "Output type, one of:\n"
6769
+ " summary - document predictions\n"
6870
+ " full - all predictions\n"
69-
+ " raw - raw response from the server"
71+
+ " raw - raw response from the server",
72+
defaultValue = "summary"
7073
)
7174
private OutputChoices outputType;
7275

@@ -110,7 +113,6 @@ public static void main(String[] args) {
110113

111114
/**
112115
* Adds all commands from CommandLineInterfaceProducts automatically.
113-
* Avoids using a Mixin, which I can't get to work.
114116
*/
115117
@CommandLine.Command(mixinStandardHelpOptions = true, description = "Auto-generated product command")
116118
public static class ProductCommandHandler implements Callable<Integer> {
@@ -139,7 +141,10 @@ public Integer call() throws Exception {
139141
}
140142
}
141143

142-
@Command(name = "custom", description = "Invokes a Custom API (API Builder only, use 'generated' for regular custom APIs)")
144+
@Command(
145+
name = "custom",
146+
description = "Invokes a Custom API (API Builder only, use 'generated' for regular custom APIs)"
147+
)
143148
void customMethod(
144149
@Option(
145150
names = {"-a", "--account"},
@@ -262,19 +267,29 @@ void generatedMethod(
262267
}
263268

264269
protected PageOptions getDefaultPageOptions() {
265-
266270
List<Integer> pageNumbers = new ArrayList<>();
267271
pageNumbers.add(0);
268272
pageNumbers.add(1);
269273
pageNumbers.add(2);
270274
pageNumbers.add(3);
271275
pageNumbers.add(4);
272-
273276
return new PageOptions(pageNumbers, PageOptionsOperation.KEEP_ONLY);
274277
}
275278

279+
private String wordsOutput(Ocr ocr) {
280+
StringBuilder output = new StringBuilder();
281+
output.append("\n#############\nDocument Text\n#############\n::\n ");
282+
output.append(
283+
Arrays.stream(ocr.toString().split(String.format("%n")))
284+
.collect(Collectors.joining(String.format("%n ")))
285+
);
286+
output.append("\n");
287+
return output.toString();
288+
}
289+
276290
@Override
277-
public <T extends Inference<?, ?>> String standardProductOutput(Class<T> productClass, File file) throws IOException {
291+
public <T extends Inference<?, ?>> String standardProductOutput(Class<T> productClass, File file)
292+
throws IOException {
278293
MindeeClient mindeeClient = new MindeeClient(apiKey);
279294
LocalInputSource inputSource = new LocalInputSource(file);
280295
PredictResponse<T> response;
@@ -285,18 +300,28 @@ protected PageOptions getDefaultPageOptions() {
285300
response = mindeeClient.parse(productClass, inputSource, predictOptions);
286301
}
287302

288-
if (outputType == OutputChoices.full) {
289-
return response.getDocument().toString();
303+
StringBuilder output = new StringBuilder();
304+
switch (outputType) {
305+
case full:
306+
output.append(response.getDocument().toString());
307+
break;
308+
case raw:
309+
output.append(response.getRawResponse());
310+
break;
311+
default:
312+
output.append(
313+
response.getDocument().getInference().getPrediction().toString()
314+
);
290315
}
291-
if (outputType == OutputChoices.raw) {
292-
return response.getRawResponse();
316+
if (words) {
317+
output.append(wordsOutput(response.getDocument().getOcr()));
293318
}
294-
Document<T> document = response.getDocument();
295-
return document.getInference().getPrediction().toString();
319+
return output.toString();
296320
}
297321

298322
@Override
299-
public <T extends Inference<?, ?>> String standardProductAsyncOutput(Class<T> productClass, File file) throws IOException, InterruptedException {
323+
public <T extends Inference<?, ?>> String standardProductAsyncOutput(Class<T> productClass, File file)
324+
throws IOException, InterruptedException {
300325
MindeeClient mindeeClient = new MindeeClient(apiKey);
301326
LocalInputSource inputSource = new LocalInputSource(file);
302327
AsyncPredictResponse<T> response;
@@ -311,13 +336,22 @@ productClass, inputSource, predictOptions, getDefaultPageOptions(), null
311336
);
312337
}
313338

314-
if (outputType == OutputChoices.full) {
315-
return response.getDocumentObj().toString();
339+
StringBuilder output = new StringBuilder();
340+
switch (outputType) {
341+
case full:
342+
output.append(response.getDocument().toString());
343+
break;
344+
case raw:
345+
output.append(response.getRawResponse());
346+
break;
347+
default:
348+
output.append(
349+
response.getDocumentObj().getInference().getPrediction().toString()
350+
);
316351
}
317-
if (outputType == OutputChoices.raw) {
318-
return response.getRawResponse();
352+
if (words) {
353+
output.append(wordsOutput(response.getDocumentObj().getOcr()));
319354
}
320-
Document<T> document = response.getDocumentObj();
321-
return document.getInference().getPrediction().toString();
355+
return output.toString();
322356
}
323357
}

0 commit comments

Comments
 (0)