Update app.py

app.py

@@ -9,6 +9,7 @@ import torch
 import easyocr
 import omegaconf
 import cv2
+import json
 
 from vietocr.vietocr.tool.predictor import Predictor
 from vietocr.vietocr.tool.config import Cfg
@@ -23,15 +24,21 @@ config['predictor']['beamsearch'] = True
 config['device'] = 'cpu' # mps
 
 recognitor = Predictor(config)
-classifier = pipeline("zero-shot-classification",
-                      model="NDugar/debertav3-mnli-snli-anli")
-def zero_shot(doc, candidates):
-    given_labels = candidates.split(", ")
-    dictionary = classifier(doc, given_labels)
-    new_dict = dict (zip (dictionary['labels'], dictionary['scores']))
-    max_label = max (new_dict, key=new_dict.get)
-    max_score = max(dictionary['scores'])
-    return max_label, max_score
+model_name = "microsoft/xdoc-base-squad2.0"
+nlp = pipeline('question-answering', model=model_name)
+def query(doc, candidates):
+    questions = candidates.split(", ")
+    result={}
+    for question in questions:
+        QA_input = {
+                    'question': question,
+                    'context': doc
+                }
+            
+        res= nlp(QA_input)
+        value = res['answer']
+        result[question]=value
+    return result
     
 def draw_boxes(image, bounds, color='yellow', width=2):
     draw = ImageDraw.Draw(image)
@@ -46,6 +53,7 @@ def inference(filepath, lang, labels):
     reader = easyocr.Reader(lang)
     bounds = reader.readtext(filepath)
     new_bounds=[]
+    text=''
     for (bbox, text, prob) in bounds:
         (tl, tr, br, bl) = bbox
         tl = (int(tl[0]), int(tl[1]))
@@ -66,15 +74,13 @@ def inference(filepath, lang, labels):
         cropped_image = img[min_y:max_y,min_x:max_x] # crop the image
         cropped_image = Image.fromarray(cropped_image)
         out = recognitor.predict(cropped_image)
-        print(out)
-        max_label, max_score = zero_shot(out, labels)
-        print(max_label)
-        print(max_score)
-        new_bounds.append((bbox,text, out, prob))
+        text = text + '\t' + out
+    result = query(text, labels)
+    jsonText = json.dumps(result)
     im = PIL.Image.open(filepath)
     draw_boxes(im, bounds)
     im.save('result.jpg')
-    return ['result.jpg', pd.DataFrame(new_bounds).iloc[: , 2:]]
+    return ['result.jpg', jsonText]
 
 title = 'EasyOCR'
 description = 'Gradio demo for EasyOCR. EasyOCR demo supports 80+ languages.To use it, simply upload your image and choose a language from the dropdown menu, or click one of the examples to load them. Read more at the links below.'
@@ -87,7 +93,7 @@ choices = [
 gr.Interface(
     inference,
     [gr.inputs.Image(type='filepath', label='Input'),gr.inputs.CheckboxGroup(choices, type="value", default=['vi'], label='language'), gr.inputs.Textbox(label='Labels')],
-    [gr.outputs.Image(type='pil', label='Output'), gr.outputs.Dataframe(type='pandas', headers=['easyOCR','vietOCR', 'confidence'])],
+    [gr.outputs.Image(type='pil', label='Output'), gr.outputs.Textbox(label='Json')],
     title=title,
     description=description,
     article=article,