Upload 3 files

app.py

@@ -0,0 +1,67 @@
+import gradio as gr
+import pandas as pd
+from models import Reciept_Analyzer
+from utils import find_product, get_info
+import os
+model = Reciept_Analyzer()
+
+sample_images = []
+for img_file in os.listdir("samples/"):
+    sample_images.append(os.path.join("samples", img_file))
+
+def predict(image):
+    results = model.forward(image)
+    return results
+
+
+
+# Thiết kế giao diện với Gradio
+def create_interface():
+    with gr.Blocks() as app:
+        gr.Markdown("# Ứng dụng phân tích hóa đơn siêu thị")
+
+        with gr.Row():
+            # Cột bên trái
+            with gr.Column():
+                gr.Markdown("### Tải lên hóa đơn hoặc chọn ảnh mẫu")
+                image_input = gr.Image(label="Ảnh hóa đơn", type="filepath")
+
+                
+
+                res = None
+                def on_image_selected(image_path):
+                    global res
+                    res = predict(image_path)
+                    final = get_info(res)
+                    print(res)
+                    return final
+
+                def handle_input(item_name):
+                    global res
+                    result = find_product(item_name, res)
+                    return result
+                
+
+                gr.Markdown("### Ảnh mẫu")
+                example = gr.Examples(
+                    inputs=image_input,
+                    examples=sample_images
+                )
+
+            # Cột bên phải
+            with gr.Column():
+                result_output = gr.Textbox(label="Kết quả phân tích")
+                image_input.change(fn=on_image_selected, inputs=image_input, outputs=result_output)
+                gr.Markdown("### Tìm kiếm thông tin item")
+                item_input = gr.Textbox(label="Tên item cần tìm")
+                output = gr.Textbox(label="Kết quả")
+
+                search_button = gr.Button("Tìm kiếm")
+                search_button.click(fn=handle_input, inputs=item_input, outputs=output)
+
+    return app
+
+
+# Chạy ứng dụng
+app = create_interface()
+app.launch()

models.py

@@ -0,0 +1,169 @@
+import os
+import torch
+import numpy as np
+from ultralytics import YOLO
+from transformers import AutoProcessor
+from transformers import AutoModelForTokenClassification
+from utils import normalize_box, unnormalize_box, draw_output, create_df
+from PIL import Image, ImageDraw
+from vietocr.tool.predictor import Predictor
+from vietocr.tool.config import Cfg
+
+class Reciept_Analyzer:
+    def __init__(self,
+                 processor_pretrained='microsoft/layoutlmv3-base',
+                 layoutlm_pretrained=os.path.join(
+                      'models', 'checkpoint'),
+                 yolo_pretrained=os.path.join(
+                      'models', 'best.pt'),
+                 vietocr_pretrained=os.path.join(
+                      'models', 'vietocr', 'vgg_seq2seq.pth')
+                 ):
+        
+        print("Initializing processor")
+        if torch.cuda.is_available():
+            print("Using GPU")
+        else:
+            print("No GPU detected, using CPU")
+
+        self.processor = AutoProcessor.from_pretrained(
+            processor_pretrained, apply_ocr=False)
+        print("Finished initializing processor")
+
+        print("Initializing LayoutLM model")
+        self.lalm_model = AutoModelForTokenClassification.from_pretrained(
+            layoutlm_pretrained)
+        print("Finished initializing LayoutLM model")
+
+        if yolo_pretrained is not None:
+            print("Initializing YOLO model")
+            self.yolo_model = YOLO(yolo_pretrained)
+            print("Finished initializing YOLO model")
+
+        print("Initializing VietOCR model")
+        config = Cfg.load_config_from_name('vgg_seq2seq')
+        config['weights'] = vietocr_pretrained
+        config['cnn']['pretrained']= False
+        config['device'] = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+        self.vietocr = Predictor(config)
+        print("Finished initializing VietOCR model")
+
+    def forward(self, img, output_path="output", is_save_cropped_img=False):
+        input_image = Image.open(img)
+
+        # detection with YOLOv8
+        bboxes = self.yolov8_det(input_image)
+        
+        # sort
+        sorted_bboxes = self.sort_bboxes(bboxes)
+        
+        # draw bbox
+        image_draw = input_image.copy()
+        self.draw_bbox(image_draw, sorted_bboxes, output_path)
+
+        # crop images
+        cropped_images, normalized_boxes = self.get_cropped_images(input_image, sorted_bboxes, is_save_cropped_img, output_path)
+        
+        # recognition with VietOCR
+        texts, mapping_bbox_texts = self.ocr(cropped_images, normalized_boxes)
+
+        # KIE with LayoutLMv3
+        pred_texts, pred_label, boxes  = self.kie(input_image, texts, normalized_boxes, mapping_bbox_texts, output_path)
+        
+        # create dataframe
+        return create_df(pred_texts, pred_label)
+        
+
+    def yolov8_det(self, img):
+        return self.yolo_model.predict(source=img, conf=0.3, iou=0.1)[0].boxes.xyxy.int()
+
+    def sort_bboxes(self, bboxes):
+        bbox_list = []
+        for box in bboxes:
+            tlx, tly, brx, bry = map(int, box)
+            bbox_list.append([tlx, tly, brx, bry])
+            bbox_list.sort(key=lambda x: (x[1], x[2]))
+        return bbox_list
+    
+    def draw_bbox(self, image_draw, bboxes, output_path):
+        # draw bbox
+        draw = ImageDraw.Draw(image_draw)
+        for box in bboxes:
+            draw.rectangle(box, outline='red', width=2)
+        image_draw.save(os.path.join(output_path, 'bbox.jpg'))
+        print(f"Exported image with bounding boxes to {os.path.join(output_path, 'bbox.jpg')}")
+
+    def get_cropped_images(self, input_image, bboxes, is_save_cropped=False, output_path="output"):
+        normalized_boxes = []
+        cropped_images = []
+
+        # OCR
+        if is_save_cropped:
+            cropped_folder = os.path.join(output_path, "cropped")
+            if not os.path.exists(cropped_folder):
+                os.makedirs(cropped_folder)
+            i = 0
+        for box in bboxes:
+            tlx, tly, brx, bry = map(int, box)
+            normalized_box = normalize_box(box, input_image.width, input_image.height)
+            normalized_boxes.append(normalized_box)
+            cropped_ = input_image.crop((tlx, tly, brx, bry))
+            if is_save_cropped:
+                cropped_.save(os.path.join(cropped_folder, f'cropped_{i}.jpg'))
+                i += 1
+            cropped_images.append(cropped_)
+
+        return cropped_images, normalized_boxes
+
+    def ocr(self, cropped_images, normalized_boxes):
+        mapping_bbox_texts = {}
+        texts = []
+        for img, normalized_box in zip(cropped_images, normalized_boxes):
+            result = self.vietocr.predict(img)
+            text = result.strip().replace('\n', ' ')
+            texts.append(text)
+            mapping_bbox_texts[','.join(map(str, normalized_box))] = text
+        
+        return texts, mapping_bbox_texts
+    
+    def kie(self, img, texts, boxes, mapping_bbox_texts, output_path):
+        encoding = self.processor(img, texts,
+                                  boxes=boxes,
+                                  return_offsets_mapping=True,
+                                  return_tensors='pt',
+                                  max_length=512,
+                                  padding='max_length')
+        offset_mapping = encoding.pop('offset_mapping')
+
+        with torch.no_grad():
+            outputs = self.lalm_model(**encoding)
+
+        id2label = self.lalm_model.config.id2label
+        logits = outputs.logits
+        token_boxes = encoding.bbox.squeeze().tolist()
+        offset_mapping = offset_mapping.squeeze().tolist()
+
+        predictions = logits.argmax(-1).squeeze().tolist()
+        is_subword = np.array(offset_mapping)[:, 0] != 0
+
+        true_predictions = []
+        true_boxes = []
+        true_texts = []
+        for idx in range(len(predictions)):
+            if not is_subword[idx] and token_boxes[idx] != [0, 0, 0, 0]:
+                true_predictions.append(id2label[predictions[idx]])
+                true_boxes.append(unnormalize_box(
+                    token_boxes[idx], img.width, img.height))
+                true_texts.append(mapping_bbox_texts.get(
+                    ','.join(map(str, token_boxes[idx])), ''))
+
+        if isinstance(output_path, str):
+            os.makedirs(output_path, exist_ok=True)
+            img_output = draw_output(
+                image=img,
+                true_predictions=true_predictions,
+                true_boxes=true_boxes
+            )
+            img_output.save(os.path.join(output_path, 'result.jpg'))
+            print(f"Exported result to {os.path.join(output_path, 'result.jpg')}")
+        return true_texts, true_predictions, true_boxes

utils.py

@@ -0,0 +1,179 @@
+import numpy as np
+from datasets import load_metric
+from PIL import ImageDraw, ImageFont
+import pandas as pd
+
+
+metric = load_metric("seqeval")
+
+
+def unnormalize_box(bbox, width, height):
+    return [
+        width * (bbox[0] / 1000),
+        height * (bbox[1] / 1000),
+        width * (bbox[2] / 1000),
+        height * (bbox[3] / 1000)
+    ]
+
+
+def normalize_box(bbox, width, height):
+    return [
+        int((bbox[0] / width) * 1000),
+        int((bbox[1] / height) * 1000),
+        int((bbox[2] / width) * 1000),
+        int((bbox[3] / height) * 1000)
+    ]
+
+
+def draw_output(image, true_predictions, true_boxes):
+    def iob_to_label(label):
+        label = label
+        if not label:
+            return 'other'
+        return label
+
+    # width, height = image.size
+
+    # predictions = logits.argmax(-1).squeeze().tolist()
+    # is_subword = np.array(offset_mapping)[:,0] != 0
+    # true_predictions = [id2label[pred] for idx, pred in enumerate(predictions) if not is_subword[idx]]
+    # true_boxes = [unnormalize_box(box, width, height) for idx, box in enumerate(token_boxes) if not is_subword[idx]]
+
+    # draw
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default()
+
+    for prediction, box in zip(true_predictions, true_boxes):
+        predicted_label = iob_to_label(prediction).lower()
+        draw.rectangle(box, outline='red')
+        draw.text((box[0] + 10, box[1] - 10),
+                  text=predicted_label, fill='red', font=font)
+
+    return image
+
+
+def create_df(true_texts,
+              true_predictions,
+              chosen_labels=['SHOP_NAME', 'ADDR', 'TITLE', 'PHONE',
+                             'PRODUCT_NAME', 'AMOUNT', 'UNIT', 'UPRICE', 'SUB_TPRICE', 'UDISCOUNT',
+                             'TAMOUNT', 'TPRICE', 'FPRICE', 'TDISCOUNT',
+                             'RECEMONEY', 'REMAMONEY',
+                             'BILLID', 'DATETIME', 'CASHIER']
+              ):
+
+    data = {'text': [], 'class_label': [], 'product_id': []}
+    product_id = -1
+    for text, prediction in zip(true_texts, true_predictions):
+        if prediction not in chosen_labels:
+            continue
+
+        if prediction == 'PRODUCT_NAME':
+            product_id += 1
+            
+
+        if prediction in ['AMOUNT', 'UNIT', 'UDISCOUNT', 'UPRICE', 'SUB_TPRICE',
+                          'UDISCOUNT', 'TAMOUNT', 'TPRICE', 'FPRICE', 'TDISCOUNT',
+                          'RECEMONEY', 'REMAMONEY']:
+            text = reformat(text)
+
+
+        if prediction in ['AMOUNT', 'SUB_TPRICE', 'UPRICE', 'PRODUCT_NAME']:
+            data['product_id'].append(product_id)
+        else:
+            data['product_id'].append('')
+
+
+        data['class_label'].append(prediction)
+        data['text'].append(text)
+
+
+    df = pd.DataFrame(data)
+
+    return df
+
+
+def reformat(text: str):
+    try:
+        text = text.replace('.', '').replace(',', '').replace(':', '').replace('/', '').replace('|', '').replace(
+            '\\', '').replace(')', '').replace('(', '').replace('-', '').replace(';', '').replace('_', '')
+        return int(text)
+    except:
+        return text
+
+def find_product(product_name, df):
+    product_name = product_name.lower()
+    product_df = df[df['class_label'] == 'PRODUCT_NAME']
+    mask = product_df['text'].str.lower().str.contains(product_name, case=False, na=False)
+    if mask.any():
+        product_id = product_df.loc[mask, 'product_id'].iloc[0]
+        product_info = df[df['product_id'] == product_id]
+        
+        prod_name = product_info.loc[product_info['class_label'] == 'PRODUCT_NAME', 'text'].iloc[0]
+            
+        try:
+            amount = product_info.loc[product_info['class_label'] == 'AMOUNT', 'text'].iloc[0]
+        except:
+            print("Error: cannot find amount")
+            amount = ''
+            
+        try:
+            uprice = product_info.loc[product_info['class_label'] == 'UPRICE', 'text'].iloc[0]
+        except:
+            print("Error: cannot find unit price")
+            uprice = ''
+            
+        try:
+            sub_tprice = product_info.loc[product_info['class_label'] == 'SUB_TPRICE', 'text'].iloc[0]
+        except:
+            print("Error: cannot find sub total price")
+            sub_tprice = ''
+            
+        #print("Sản phẩm: ", product_info.loc[product_info['class_label'] == 'PRODUCT_NAME', 'text'].iloc[0])
+        #print("Số lượng: ", product_info.loc[product_info['class_label'] == 'AMOUNT', 'text'].iloc[0])
+        #print("Đơn giá: ", product_info.loc[product_info['class_label'] == 'UPRICE', 'text'].iloc[0])
+        #print("Thành tiền: ", product_info.loc[product_info['class_label'] == 'SUB_TPRICE', 'text'].iloc[0])
+        return f"Sản phẩm: {prod_name}\n Số lượng: {amount}\n Đơn giá: {uprice}\n Thành tiền: {sub_tprice}"
+    else:
+        #print("Không tìm thấy item nào phù hợp.")
+        return "Không tìm thấy item nào phù hợp."
+    #return result = product_df['text'].str.contains(product_name, case=False, na=False).any()
+    #return product_df[product_df['text'].str.contains(product_name, case=False, na=False)]
+
+
+def get_info(df):
+    try:
+        shop_name = df.loc[df['class_label'] == 'SHOP_NAME', 'text'].iloc[0]
+    except:
+        print("Error: cannot find shop name")
+        shop_name = ''
+    print("Tên siêu thị: ", shop_name)
+
+    try:
+        addr = df.loc[df['class_label'] == 'ADDR', 'text'].iloc[0]
+    except:
+        print("Error: cannot find address")
+        addr = ''
+    print("Địa chỉ: ", addr)
+
+    try:
+        bill_id = df.loc[df['class_label'] == 'BILLID', 'text'].iloc[0]
+    except:
+        print("Error: cannot find bill id")
+        bill_id = ''
+    print("ID hóa đơn: ", bill_id)
+
+    try:
+        date_time = df.loc[df['class_label'] == 'DATETIME', 'text'].iloc[0]
+    except:
+        print("Error: cannot find date and time")
+        date_time = ''
+    print("Ngày: ", date_time)
+
+    try:
+        cashier = df.loc[df['class_label'] == 'CASHIER', 'text'].iloc[0]
+    except:
+        print("Error: cannot find cashier")
+        cashier = ''
+    print("Nhân viên: ", cashier)
+
+    return f"Tên siêu thị: {shop_name}\n Địa chỉ: {addr}\n ID hóa đơn: {bill_id}\n Ngày: {date_time}\n Nhân viên: {cashier}\n"