Update app.py

app.py

@@ -1,4 +1,6 @@
 import pandas as pd
+from transformers import pipeline
+
 import PIL
 from PIL import Image
 from PIL import ImageDraw
@@ -21,7 +23,16 @@ 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
+    
 def draw_boxes(image, bounds, color='yellow', width=2):
     draw = ImageDraw.Draw(image)
     for bound in bounds:
@@ -29,7 +40,7 @@ def draw_boxes(image, bounds, color='yellow', width=2):
         draw.line([*p0, *p1, *p2, *p3, *p0], fill=color, width=width)
     return image
 
-def inference(filepath, lang):
+def inference(filepath, lang, labels):
     img = cv2.imread(filepath)
     width, height, _ = img.shape
     reader = easyocr.Reader(lang)
@@ -56,6 +67,9 @@ def inference(filepath, lang):
         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))
     im = PIL.Image.open(filepath)
     draw_boxes(im, bounds)
@@ -72,7 +86,7 @@ choices = [
 ]
 gr.Interface(
     inference,
-    [gr.inputs.Image(type='filepath', label='Input'),gr.inputs.CheckboxGroup(choices, type="value", default=['vi'], label='language')],
+    [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'])],
     title=title,
     description=description,