Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import spacy | |
| import json | |
| import os | |
| import re | |
| try: | |
| nlp = spacy.load("./model") | |
| model_loaded = True | |
| except: | |
| model_loaded = False | |
| corrections = [] | |
| if os.path.exists("corrections.json"): | |
| with open("corrections.json", "r") as f: | |
| corrections = json.load(f) | |
| def extract_postcode(text): | |
| if not text: | |
| return "" | |
| pattern = r'([A-Z]{1,2}\d{1,2}[A-Z]?\s*\d[A-Z]{2})' | |
| match = re.search(pattern, text.upper()) | |
| if match: | |
| postcode = match.group(1) | |
| if ' ' not in postcode: | |
| postcode = re.sub(r'(\d)([A-Z]{2})$', r'\1 \2', postcode) | |
| return postcode | |
| return "" | |
| def extract_from_text(ocr_text): | |
| if not model_loaded: | |
| return tuple(["Model not loaded"] + [""]*23 + [False]*6 + [1, 1]) | |
| if not ocr_text.strip(): | |
| return tuple(["Paste OCR text"] + [""]*23 + [False]*6 + [1, 1]) | |
| doc = nlp(ocr_text) | |
| price, pickups, deliveries, restaurants = "", [], [], [] | |
| for ent in doc.ents: | |
| if ent.label_ == "PRICE": price = ent.text | |
| elif ent.label_ == "RESTAURANT": restaurants.append(ent.text) | |
| elif ent.label_ == "PICKUP": pickups.append(ent.text) | |
| elif ent.label_ == "DELIVERY": deliveries.append(ent.text) | |
| is_stacked = len(deliveries) > 1 or bool(re.search(r'Delivery\s*\((\d+)\)', ocr_text)) | |
| currency = "GBP" | |
| if "$" in ocr_text: currency = "USD" | |
| elif "AED" in ocr_text: currency = "AED" | |
| order_indicator = "" | |
| if re.search(r'Delivery\s*\(\d+\)', ocr_text): | |
| match = re.search(r'Delivery\s*\(\d+\)', ocr_text) | |
| order_indicator = match.group(0) if match else "" | |
| elif "Delivery" in ocr_text: order_indicator = "Delivery" | |
| elif re.search(r'\d+\s*stops?', ocr_text): | |
| match = re.search(r'\d+\s*stops?', ocr_text) | |
| order_indicator = match.group(0) if match else "" | |
| elif re.search(r'\d+\s*orders?', ocr_text): | |
| match = re.search(r'\d+\s*orders?', ocr_text) | |
| order_indicator = match.group(0) if match else "" | |
| button_text = "" | |
| if "Confirm" in ocr_text: button_text = "Confirm" | |
| elif "Accept offer" in ocr_text: button_text = "Accept offer" | |
| elif "Accept and go" in ocr_text: button_text = "Accept and go" | |
| distance = "" | |
| dist_match = re.search(r'(\d+\.?\d*)\s*mi', ocr_text) | |
| if dist_match: distance = dist_match.group(0) | |
| time_est = "" | |
| time_match = re.search(r'(\d+)\s*min', ocr_text) | |
| if time_match: time_est = time_match.group(0) | |
| has_exclusive = "Exclusive" in ocr_text or "exclusive" in ocr_text | |
| has_bonus = "bonus" in ocr_text.lower() or "tip" in ocr_text.lower() | |
| same_dropoff = "same drop" in ocr_text.lower() | |
| customer_verify = "verification" in ocr_text.lower() | |
| delivery_code = "delivery code" in ocr_text.lower() | |
| p1_pc = extract_postcode(pickups[0]) if len(pickups) > 0 else "" | |
| p1_addr = pickups[0] if len(pickups) > 0 else "" | |
| p1_biz_area = restaurants[0] if len(restaurants) > 0 else "" | |
| p1_biz = restaurants[0].split()[0] if len(restaurants) > 0 else "" | |
| p2_pc = extract_postcode(pickups[1]) if len(pickups) > 1 else "" | |
| p2_addr = pickups[1] if len(pickups) > 1 else "" | |
| p2_biz_area = restaurants[1] if len(restaurants) > 1 else "" | |
| p2_biz = restaurants[1].split()[0] if len(restaurants) > 1 else "" | |
| d1_pc = extract_postcode(deliveries[0]) if len(deliveries) > 0 else "" | |
| d1_addr = deliveries[0] if len(deliveries) > 0 else "" | |
| d1_name = "" | |
| d2_pc = extract_postcode(deliveries[1]) if len(deliveries) > 1 else "" | |
| d2_addr = deliveries[1] if len(deliveries) > 1 else "" | |
| d2_name = "" | |
| d3_pc = extract_postcode(deliveries[2]) if len(deliveries) > 2 else "" | |
| d3_addr = deliveries[2] if len(deliveries) > 2 else "" | |
| d3_name = "" | |
| pickup_count = len(pickups) if len(pickups) > 0 else 1 | |
| delivery_count = len(deliveries) if len(deliveries) > 0 else 1 | |
| status = "Found " + str(pickup_count) + " pickup(s), " + str(delivery_count) + " delivery(s)" | |
| return (status, price, currency, order_indicator, button_text, distance, time_est, | |
| p1_pc, p1_addr, p1_biz_area, p1_biz, | |
| p2_pc, p2_addr, p2_biz_area, p2_biz, | |
| d1_pc, d1_addr, d1_name, | |
| d2_pc, d2_addr, d2_name, | |
| d3_pc, d3_addr, d3_name, | |
| is_stacked, has_exclusive, has_bonus, same_dropoff, customer_verify, delivery_code, | |
| pickup_count, delivery_count) | |
| def save_data(ocr_input, price, currency, order_indicator, button_text, distance, time_est, | |
| p1_pc, p1_addr, p1_biz_area, p1_biz, | |
| p2_pc, p2_addr, p2_biz_area, p2_biz, | |
| d1_pc, d1_addr, d1_name, | |
| d2_pc, d2_addr, d2_name, | |
| d3_pc, d3_addr, d3_name, | |
| is_stacked, has_exclusive, has_bonus, same_dropoff, customer_verify, delivery_code, | |
| pickup_count, delivery_count, status_type): | |
| pickups = [] | |
| if p1_pc or p1_addr or p1_biz_area or p1_biz: | |
| pickups.append({ | |
| "pickup_number": 1, | |
| "postcode": p1_pc, | |
| "full_address": p1_addr, | |
| "business_area": p1_biz_area, | |
| "business_only": p1_biz | |
| }) | |
| if p2_pc or p2_addr or p2_biz_area or p2_biz: | |
| pickups.append({ | |
| "pickup_number": 2, | |
| "postcode": p2_pc, | |
| "full_address": p2_addr, | |
| "business_area": p2_biz_area, | |
| "business_only": p2_biz | |
| }) | |
| dels = [] | |
| if d1_pc or d1_addr: | |
| dels.append({ | |
| "delivery_number": 1, | |
| "postcode": d1_pc, | |
| "full_address": d1_addr, | |
| "customer_name": d1_name | |
| }) | |
| if d2_pc or d2_addr: | |
| dels.append({ | |
| "delivery_number": 2, | |
| "postcode": d2_pc, | |
| "full_address": d2_addr, | |
| "customer_name": d2_name | |
| }) | |
| if d3_pc or d3_addr: | |
| dels.append({ | |
| "delivery_number": 3, | |
| "postcode": d3_pc, | |
| "full_address": d3_addr, | |
| "customer_name": d3_name | |
| }) | |
| corrections.append({ | |
| "ocr": ocr_input, | |
| "data": { | |
| "basic_info": { | |
| "price": price, | |
| "currency": currency, | |
| "is_stacked": is_stacked, | |
| "pickup_count": pickup_count, | |
| "delivery_count": delivery_count | |
| }, | |
| "platform_detection": { | |
| "order_indicator": order_indicator, | |
| "button_text": button_text, | |
| "distance": distance, | |
| "time_estimate": time_est, | |
| "has_exclusive": has_exclusive, | |
| "has_bonus": has_bonus | |
| }, | |
| "pickups": pickups, | |
| "deliveries": dels, | |
| "special_flags": { | |
| "same_dropoff": same_dropoff, | |
| "customer_verification": customer_verify, | |
| "delivery_code_required": delivery_code | |
| } | |
| }, | |
| "status": status_type | |
| }) | |
| with open("corrections.json", "w") as f: | |
| json.dump(corrections, f, indent=2) | |
| acc = sum(1 for c in corrections if c["status"] == "accepted") | |
| return "Saved! Total: " + str(len(corrections)) + " | Correct: " + str(acc) + " | Fixed: " + str(len(corrections)-acc) | |
| def accept(*args): | |
| return save_data(*args, "accepted") | |
| def correct(*args): | |
| return save_data(*args, "corrected") | |
| def get_file(): | |
| return "corrections.json" if os.path.exists("corrections.json") else None | |
| with gr.Blocks(title="StackAdvisor Tester") as demo: | |
| gr.Markdown("# StackAdvisor AI Tester") | |
| with gr.Row(): | |
| ocr_input = gr.Textbox(label="Paste OCR Text", lines=12) | |
| result = gr.Textbox(label="Status", lines=2) | |
| extract_btn = gr.Button("Extract", variant="primary") | |
| gr.Markdown("## BASIC INFO") | |
| with gr.Row(): | |
| price = gr.Textbox(label="Price") | |
| currency = gr.Dropdown(label="Currency", choices=["GBP", "USD", "AED"], value="GBP") | |
| is_stacked = gr.Checkbox(label="Stacked?") | |
| with gr.Row(): | |
| pickup_count = gr.Number(label="Pickup Count", value=1, minimum=1, maximum=2) | |
| delivery_count = gr.Number(label="Delivery Count", value=1, minimum=1, maximum=3) | |
| gr.Markdown("## PLATFORM DETECTION") | |
| with gr.Row(): | |
| order_indicator = gr.Textbox(label="Order Indicator") | |
| button_text = gr.Textbox(label="Button Text") | |
| with gr.Row(): | |
| distance = gr.Textbox(label="Distance") | |
| time_est = gr.Textbox(label="Time Estimate") | |
| with gr.Row(): | |
| has_exclusive = gr.Checkbox(label="Exclusive Badge") | |
| has_bonus = gr.Checkbox(label="Bonus/Tip") | |
| gr.Markdown("## PICKUP 1") | |
| with gr.Row(): | |
| p1_pc = gr.Textbox(label="Postcode") | |
| p1_addr = gr.Textbox(label="Full Address") | |
| p1_biz_area = gr.Textbox(label="Business+Area") | |
| p1_biz = gr.Textbox(label="Business Only") | |
| gr.Markdown("## PICKUP 2") | |
| with gr.Row(): | |
| p2_pc = gr.Textbox(label="Postcode") | |
| p2_addr = gr.Textbox(label="Full Address") | |
| p2_biz_area = gr.Textbox(label="Business+Area") | |
| p2_biz = gr.Textbox(label="Business Only") | |
| gr.Markdown("## DELIVERY 1") | |
| with gr.Row(): | |
| d1_pc = gr.Textbox(label="Postcode") | |
| d1_addr = gr.Textbox(label="Full Address") | |
| d1_name = gr.Textbox(label="Customer Name") | |
| gr.Markdown("## DELIVERY 2") | |
| with gr.Row(): | |
| d2_pc = gr.Textbox(label="Postcode") | |
| d2_addr = gr.Textbox(label="Full Address") | |
| d2_name = gr.Textbox(label="Customer Name") | |
| gr.Markdown("## DELIVERY 3") | |
| with gr.Row(): | |
| d3_pc = gr.Textbox(label="Postcode") | |
| d3_addr = gr.Textbox(label="Full Address") | |
| d3_name = gr.Textbox(label="Customer Name") | |
| gr.Markdown("## SPECIAL FLAGS") | |
| with gr.Row(): | |
| same_dropoff = gr.Checkbox(label="Same Drop-off") | |
| customer_verify = gr.Checkbox(label="Customer Verification") | |
| delivery_code = gr.Checkbox(label="Delivery Code Required") | |
| with gr.Row(): | |
| accept_btn = gr.Button("Accept - AI Correct", variant="primary") | |
| correct_btn = gr.Button("Save My Correction", variant="secondary") | |
| status = gr.Textbox(label="Save Status") | |
| download_btn = gr.Button("Download Corrections") | |
| download_file = gr.File() | |
| all_inputs = [ocr_input, price, currency, order_indicator, button_text, distance, time_est, | |
| p1_pc, p1_addr, p1_biz_area, p1_biz, | |
| p2_pc, p2_addr, p2_biz_area, p2_biz, | |
| d1_pc, d1_addr, d1_name, | |
| d2_pc, d2_addr, d2_name, | |
| d3_pc, d3_addr, d3_name, | |
| is_stacked, has_exclusive, has_bonus, same_dropoff, customer_verify, delivery_code, | |
| pickup_count, delivery_count] | |
| all_outputs = [result, price, currency, order_indicator, button_text, distance, time_est, | |
| p1_pc, p1_addr, p1_biz_area, p1_biz, | |
| p2_pc, p2_addr, p2_biz_area, p2_biz, | |
| d1_pc, d1_addr, d1_name, | |
| d2_pc, d2_addr, d2_name, | |
| d3_pc, d3_addr, d3_name, | |
| is_stacked, has_exclusive, has_bonus, same_dropoff, customer_verify, delivery_code, | |
| pickup_count, delivery_count] | |
| extract_btn.click(extract_from_text, inputs=[ocr_input], outputs=all_outputs) | |
| accept_btn.click(accept, inputs=all_inputs, outputs=[status]) | |
| correct_btn.click(correct, inputs=all_inputs, outputs=[status]) | |
| download_btn.click(get_file, outputs=[download_file]) | |
| demo.launch() |