add raw text output

app.py

@@ -8,7 +8,7 @@ import gradio as gr
 
 # import sys
 # sys.path.insert(0, './')
-from test import create_letr, draw_fig
+from test import create_letr, get_lines_and_draw
 from models.preprocessing import *
 from models.misc import nested_tensor_from_tensor_list
 
@@ -57,9 +57,9 @@ def predict(inp, size, model_name):
         else:
             outputs = model(inputs)[0]
 
-    draw_fig(image, outputs, orig_size)
+    lines = get_lines_and_draw(image, outputs, orig_size)
 
-    return image
+    return image, str(lines)
 
 
 inputs = [
@@ -67,7 +67,10 @@ inputs = [
     gr.inputs.Radio(["256", "512", "1100"]),
     gr.inputs.Radio(["resnet50", "resnet101"]),
 ]
-outputs = gr.outputs.Image()
+outputs = [
+    gr.outputs.Image(),
+    gr.outputs.Textbox()
+]
 gr.Interface(
     fn=predict,
     inputs=inputs,

test.py

@@ -19,7 +19,7 @@ def create_letr(path):
     model.eval()
     return model
 
-def draw_fig(image, outputs, orig_size):
+def get_lines_and_draw(image, outputs, orig_size):
     # find lines
     out_logits, out_line = outputs['pred_logits'], outputs['pred_lines']
     prob = F.softmax(out_logits, -1)
@@ -42,6 +42,7 @@ def draw_fig(image, outputs, orig_size):
         for tp_id, line in enumerate(lines):
             y1, x1, y2, x2 = line
             draw.line((x1, y1, x2, y2), fill=500)
+    return lines
 
 if __name__ == '__main__':
     model = create_letr('resnet50/checkpoint0024.pth')
@@ -62,6 +63,6 @@ if __name__ == '__main__':
 
     with torch.no_grad():
         outputs = model(inputs)[0]
-    draw_fig(image, outputs, orig_size)
+    lines = get_lines_and_draw(image, outputs, orig_size)
 
     image.save('output.png')