Sivakkanth commited on
Commit
4a4c982
·
1 Parent(s): 4029c92

Updated receipt details extract function

Browse files
Files changed (3) hide show
  1. app.py +65 -0
  2. model/best.pt +3 -0
  3. requirements.txt +6 -0
app.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import cv2
3
+ from ultralytics import YOLO
4
+ import easyocr
5
+ import numpy as np
6
+
7
+ # Load YOLO model (replace with your custom weights path if needed)
8
+ model = YOLO("model/best.pt")
9
+
10
+ # Initialize OCR
11
+ reader = easyocr.Reader(['en'])
12
+
13
+ # Class names
14
+ class_names = ["Merchant","date","total","no","item"]
15
+
16
+ def extract_receipt(image):
17
+ # Convert from PIL to OpenCV
18
+ img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
19
+
20
+ results = model(img)[0]
21
+
22
+ output = {
23
+ "items": [],
24
+ "name": "",
25
+ "total": "",
26
+ "date": "",
27
+ "time": "",
28
+ "discount": "",
29
+ "tax": ""
30
+ }
31
+
32
+ for box, cls_id, conf in zip(results.boxes.xyxy, results.boxes.cls, results.boxes.conf):
33
+ x1, y1, x2, y2 = [int(i) for i in box]
34
+ cls_id = int(cls_id)
35
+ cls_name = class_names[cls_id]
36
+ crop = img[y1:y2, x1:x2]
37
+ text_result = reader.readtext(crop)
38
+ text = " ".join([t[1] for t in text_result])
39
+
40
+ if cls_name == "Merchant":
41
+ output["name"] = text
42
+ elif cls_name == "date":
43
+ output["date"] = text
44
+ elif cls_name == "total":
45
+ output["total"] = text
46
+ elif cls_name == "no":
47
+ output["time"] = text
48
+ elif cls_name == "item":
49
+ parts = text.rsplit(" ", 1)
50
+ if len(parts) == 2 and parts[1].replace(".","").isdigit():
51
+ output["items"].append({"product": parts[0], "price": parts[1]})
52
+ else:
53
+ output["items"].append({"product": text, "price": ""})
54
+ return output
55
+
56
+ # Gradio interface
57
+ iface = gr.Interface(
58
+ fn=extract_receipt,
59
+ inputs=gr.Image(type="pil"),
60
+ outputs=gr.JSON(),
61
+ title="Receipt Extractor",
62
+ description="Upload a receipt image to extract merchant, date, total, and items."
63
+ )
64
+
65
+ iface.launch()
model/best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:81a4e87311bbd29b5c4e7d76445323cbcca276f788ff6900e4957e1ee90bae42
3
+ size 6250602
requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ ultralytics
4
+ opencv-python-headless
5
+ easyocr
6
+ gradio