First update

app.py

@@ -4,54 +4,103 @@ from PIL import Image
 import gradio as gr
 import numpy as np
 import cv2
+from gliner import GLiNER
+
+# Initialize GLiNER model
+gliner_model = GLiNER.from_pretrained("urchade/gliner_large-v2.1")
+
+# Entity labels including website
+labels = ["person name", "company name", "job title", "phone", "email", "address", "website"]
 
-# 获取随机的颜色
 def get_random_color():
     c = tuple(np.random.randint(0, 256, 3).tolist())
     return c
 
-# 绘制ocr识别结果
 def draw_ocr_bbox(image, boxes, colors):
-    print(colors)
-    box_num = len(boxes)
-    for i in range(box_num):
-        box = np.reshape(np.array(boxes[i]), [-1, 1, 2]).astype(np.int64)
-        image = cv2.polylines(np.array(image), [box], True, colors[i], 2)
+    valid_boxes = []
+    valid_colors = []
+    for box, color in zip(boxes, colors):
+        if len(box) > 0:  # Only draw valid boxes
+            valid_boxes.append(box)
+            valid_colors.append(color)
+    
+    for box, color in zip(valid_boxes, valid_colors):
+        box = np.array(box).reshape(-1, 1, 2).astype(np.int64)
+        image = cv2.polylines(np.array(image), [box], True, color, 2)
     return image
 
-# torch.hub.download_url_to_file('https://i.imgur.com/aqMBT0i.jpg', 'example.jpg')
-
 def inference(img: Image.Image, lang, confidence):
+    # Initialize PaddleOCR
     ocr = PaddleOCR(use_angle_cls=True, lang=lang, use_gpu=False,
                     det_model_dir=f'./models/det/{lang}',
                     cls_model_dir=f'./models/cls/{lang}',
                     rec_model_dir=f'./models/rec/{lang}')
-    # img_path = img.name
+    
+    # Process image
     img2np = np.array(img)
-    result = ocr.ocr(img2np, cls=True)[0]
-    # rgb
+    ocr_result = ocr.ocr(img2np, cls=True)[0]
+    
+    # Original OCR processing
+    ocr_items = []
+    if ocr_result:
+        boxes = [line[0] for line in ocr_result]
+        txts = [line[1][0] for line in ocr_result]
+        scores = [line[1][1] for line in ocr_result]
+        
+        ocr_items = [
+            {'boxes': box, 'txt': txt, 'score': score, '_c': get_random_color()}
+            for box, txt, score in zip(boxes, txts, scores)
+            if score > confidence
+        ]
+    
+    # GLiNER Entity Extraction
+    combined_text = " ".join([item['txt'] for item in ocr_items])
+    gliner_entities = gliner_model.predict_entities(combined_text, labels, threshold=0.3)
+    
+    # Add GLiNER entities (without boxes)
+    gliner_items = [
+        {'boxes': [], 'txt': f"{ent['text']} ({ent['label']})", 'score': 1.0, '_c': get_random_color()}
+        for ent in gliner_entities
+    ]
+    
+    # QR Code Detection
+    qr_items = []
+    qr_detector = cv2.QRCodeDetector()
+    retval, decoded_info, points, _ = qr_detector.detectAndDecodeMulti(img2np)
+    
+    if retval:
+        for i, url in enumerate(decoded_info):
+            if url:
+                qr_box = points[i].reshape(-1, 2).tolist()
+                qr_items.append({
+                    'boxes': qr_box,
+                    'txt': url,
+                    'score': 1.0,
+                    '_c': get_random_color()
+                })
+    
+    # Combine all results
+    final_result = ocr_items + gliner_items + qr_items
+    
+    # Prepare output
     image = img.convert('RGB')
-    boxes = [line[0] for line in result]
-    txts = [line[1][0] for line in result]
-    scores = [line[1][1] for line in result]
+    image_with_boxes = draw_ocr_bbox(image, 
+                                   [item['boxes'] for item in final_result],
+                                   [item['_c'] for item in final_result])
     
-    # 识别结果
-    final_result = [dict(boxes=box, txt=txt, score=score, _c=get_random_color()) for box, txt, score in zip(boxes, txts, scores)]
-    # 过滤 score < 0.5 的
-    final_result = [item for item in final_result if item['score'] > confidence]
-
-    im_show = draw_ocr_bbox(image, [item['boxes'] for item in final_result], [item['_c'] for item in final_result])
-    im_show = Image.fromarray(im_show)
-    data = [[json.dumps(item['boxes']), round(item['score'], 3), item['txt']] for item in final_result]
-    return im_show, data
+    data = [
+        [json.dumps(item['boxes']), round(item['score'], 3), item['txt']]
+        for item in final_result
+    ]
+    
+    return Image.fromarray(image_with_boxes), data
 
-title = 'PaddleOCR'
-description = 'Gradio demo for PaddleOCR.'
+title = 'Enhanced Business Card Scanner'
+description = 'Combines OCR, entity recognition, and QR scanning'
 
 examples = [
-    ['example_imgs/example.jpg','en', 0.5],
-    ['example_imgs/ch.jpg','ch', 0.7],
-    ['example_imgs/demo003.jpeg','en', 0.7],
+    ['example_imgs/example.jpg', 'en', 0.5],
+    ['example_imgs/demo003.jpeg', 'en', 0.7],
 ]
 
 css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
@@ -59,17 +108,15 @@ css = ".output_image, .input_image {height: 40rem !important; width: 100% !impor
 if __name__ == '__main__':
     demo = gr.Interface(
         inference,
-        [gr.Image(type='pil', label='Input'),
-        gr.Dropdown(choices=['ch', 'en', 'fr', 'german', 'korean', 'japan'], value='ch', label='language'),
-        gr.Slider(0.1, 1, 0.5, step=0.1, label='confidence_threshold')
+        [
+            gr.Image(type='pil', label='Input'),
+            gr.Dropdown(choices=['en', 'fr', 'german', 'korean', 'japan'], value='en', label='Language'),
+            gr.Slider(0.1, 1, 0.5, step=0.1, label='Confidence Threshold')
         ],
-        # 输出
-        [gr.Image(type='pil', label='Output'), gr.Dataframe(headers=[ 'bbox', 'score', 'text'], label='Result')],
+        [gr.Image(type='pil', label='Output'), gr.Dataframe(headers=['bbox', 'score', 'text'], label='Results')],
         title=title,
         description=description,
         examples=examples,
-        css=css,
-        cache_examples=True  # 添加缓存选项
+        css=css
     )
-    demo.queue(max_size=10)
-    demo.launch()
+    demo.launch()