return json

app.py

@@ -27,7 +27,7 @@ def run_inference(model_path, images_path):
         inference_batch = prepare_batch_for_inference(images_path)
         context = {"model_dir": model_path}
         print('handle....')
-        handle(inference_batch,context)
+        return handle(inference_batch,context)
     except Exception as err:        
         print('err...', err)
 
@@ -44,15 +44,21 @@ def run(img_path):
     print('img path: ', img_path)
     model_path = get_model()
     model_path = pretrained_model(model_path)
-    run_inference(model_path, [img_path])
+    results = run_inference(model_path, [img_path])
 
-    img_name = os.path.basename(img_path)
-    ext =  os.path.splitext(img_name)[1]
-    img_name = img_name[:img_name.find('.')]    
+    if len(results) == 0:
+        return Image.open(img_path), {'status': 'error'}
+    
+    # img_name = os.path.basename(img_path)
+    # ext =  os.path.splitext(img_name)[1]
+    # img_name = img_name[:img_name.find('.')]    
+
+    # return Image.open(os.path.join('tmp', f'{img_name}_inference.jpg')) 
 
-    return Image.open(os.path.join('tmp', f'{img_name}_inference.jpg')) 
+    print(results[0])
+    return Image.open(results[0]['image_path']), results[0]
 
 
 gr.Markdown('Upload Foto Wajah Kamu (Pastikan hanya terdapat SATU wajah pada)')
-iface = gr.Interface(fn=run, inputs=gr.Image(type="filepath"), outputs="image")
+iface = gr.Interface(fn=run, inputs=gr.Image(type="filepath"), outputs=["image","json"])
 iface.launch()

inference/annotate_image.py

@@ -34,8 +34,10 @@ def annotate_image(image_path, annotation_object):
   draw = ImageDraw.Draw(overlay)
   font = ImageFont.load_default()
 
+  print('annotation_object ---> ',annotation_object)
   predictions = [span['label'] for span in annotation_object['output']]
   boxes = [span['words'][0]['box'] for span in annotation_object['output']]
+  receipt_info = [{span['label']: span['text']} for span in annotation_object['output']]
   for prediction, box in zip(predictions, boxes):
       draw.rectangle(box, outline=label2color[prediction],
                      width=3, fill=label2color[prediction]+(int(255*0.33),))
@@ -47,4 +49,6 @@ def annotate_image(image_path, annotation_object):
 
   image_name = os.path.basename(image_path)
   image_name = image_name[:image_name.find('.')]
-  img.save(f'tmp/{image_name}_inference.jpg')
+  image_inf_path = f'tmp/{image_name}_inference.jpg'
+  img.save(image_inf_path)
+  return {'image_path': image_inf_path, 'data': receipt_info}

inference/inference_handler.py

@@ -178,8 +178,13 @@ class ModelHandler(object):
             inf_out.write(inference_out)
         inference_out_list = json.loads(inference_out)
         flattened_output_list = get_flattened_output(inference_out_list)
+
+        result = []
         for i, flattened_output in enumerate(flattened_output_list):
-            annotate_image(data['image_path'][i], flattened_output)
+            res = annotate_image(data['image_path'][i], flattened_output)
+            result.append(res)
+        
+        return result
             
 
 

inference/ocr.py

@@ -56,5 +56,5 @@ def prepare_batch_for_inference(image_paths):
       "bboxes": boxes_lists,
       "words": word_lists
   }
-  print('inference_batch:', inference_batch)
+#   print('inference_batch:', inference_batch)
   return inference_batch