Spaces:
No application file
No application file
| import urllib | |
| import pandas as pd | |
| import numpy as np | |
| from pdf2image import convert_from_path | |
| import cv2 | |
| import pytesseract | |
| import requests | |
| import json | |
| import re | |
| from roboflow import Roboflow | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| import pypdfium2 as pdfium | |
| search_google = True | |
| GOOGLE_API_KEY = "AIzaSyCD5eIu3N2uD6Vf5GE0VLBpkPsJ3mEZgA0" | |
| # GOOGLE_CSE_ID = "b536303d151b94476" | |
| # GOOGLE_API_KEY = "AIzaSyBUvPKAZa2TddKYDJI9t0ac2A_JJqdgWHo" | |
| # GOOGLE_CSE_ID = "f3446c60c52cd4dc2" | |
| # GOOGLE_API_KEY = "AIzaSyCoBYZK4khX_W5eW-AqEXKOclUIer6CiC8" | |
| # GOOGLE_CSE_ID = "c3ba01991c1a14c7d" | |
| GOOGLE_CSE_ID = "b536303d151b94476" | |
| mode = 'local' | |
| # mode = 'rbf' | |
| if mode == 'local': | |
| model = YOLO('best-flyers-yolov8-1000img-v10.pt') | |
| else: | |
| rf = Roboflow(api_key="TN586SQsdwjUUvMN3H8e") | |
| project = rf.workspace().project("smartflyers") | |
| model = project.version(9).model | |
| # def process(selected_shop): | |
| # print(selected_shop) | |
| # return selected_shop | |
| hotspots_sku_links = [] | |
| shop = None | |
| shops = { | |
| 'Latulippe': { | |
| 'url': 'https://latulippe.com/fr/recherche/?keywords=', | |
| 'regex': r'N° (M\d+)', | |
| 'multi_items': False, | |
| }, | |
| 'Decathlon': { | |
| 'url': 'https://www.decathlon.fr/search?Ntt=', | |
| 'regex': r'Réf\. ?: ?([\d, ]+)', | |
| 'multi_items': True, | |
| 'brands': [ | |
| "Allsix", "Aptonia", "Artengo", "Atorka", "B'Twin", "Caperlan", "Domyos", "Dreamscape", | |
| "FLX", "Forclaz", "Fouganza", "Geologic", "Geonaute", "Inesis", "Imviso", "Inovik", | |
| "Itiwit", "KOODZA", "KOROK", "Kalenji", "Kimjaly", "Kipsta", "Kuikma", "Nabaiji", | |
| "Newfeel", "Offload", "Olaian", "Opfeel", "Orao", "Oroks", "Outshock", "Oxelo", | |
| "Perfly", "Pongori", "Quechua", "Rockrider", "Sandever", "Simond", "Solognac", | |
| "Subea", "Tarmak", "Triban", "Tribord", "Urball", "Van Rysel", "Wed'Ze", "Decathlon" | |
| ] | |
| }, | |
| 'Canadian Tire': { | |
| 'url': 'https://www.canadiantire.ca/en/search-results.html?q=', | |
| 'regex': r'(\d{2,3}-\d{4}-\d(?:/\d{4}-\d)?)', | |
| 'multi_items': True, | |
| }, | |
| 'Walmart': { | |
| # 'url': 'https://www.walmart.ca/en/search?q=', | |
| 'url': '', | |
| 'regex': r'N°(\d+)', | |
| 'multi_items': False, | |
| }, | |
| 'Colruyt': { | |
| 'url': 'https://www.colruyt.be/fr/producten?searchTerm=e', | |
| 'regex': r'', | |
| 'multi_items': False, | |
| }, | |
| # Add more shops and their respective information here | |
| } | |
| regex_price = r'(\d+\.?\d*) ?([€$])|([€$]) ?(\d+\.?\d*)' | |
| def pdf_to_image(pdf): | |
| pdf = pdfium.PdfDocument(pdf) | |
| n_pages = len(pdf) # get the number of pages in the document | |
| print(f"Number of pages: {n_pages}") | |
| assert n_pages > 0, "The PDF document is empty" | |
| assert n_pages == 1, "The PDF document has more than one page" | |
| page = pdf.get_page(0) | |
| bitmap = page.render( | |
| scale=5, # 72dpi resolution | |
| rotation=0, # no additional rotation | |
| # ... further rendering options | |
| ) | |
| pil_image = bitmap.to_pil() | |
| pil_image.save('test.png') | |
| width_img, height_img = pil_image.size | |
| print("width image", width_img, "height image", height_img) | |
| width_pdf, height_pdf = page.get_size() | |
| print("width pdf ", width_pdf, "height pdf ", height_pdf) | |
| #assert int(width_pdf) == width_img, "The width of the image is not the same as the width of the page" | |
| #assert int(height_pdf) == height_img, "The height of the image is not the same as the height of the page" | |
| return pil_image, width_pdf, height_pdf, page | |
| def app(pdf, selected_shop): | |
| if pdf is None: | |
| raise gr.Error("No file uploaded. Please upload PDF.") | |
| # Check if the uploaded file is a PDF | |
| if selected_shop is None: | |
| raise gr.Error("Please select a shop.") | |
| try: | |
| image_annotated = section(pdf) | |
| pass | |
| except Exception as e: | |
| return gr.Error(f"An error occurred: {str(e)}") | |
| return image_annotated, hotspots_sku_links[selected_shop] | |
| def build_request(query, start=1, num=10, date_restrict=None, **params): | |
| request = { | |
| 'key': GOOGLE_API_KEY, | |
| 'q': query, | |
| 'cx': GOOGLE_CSE_ID, | |
| # 'start': start, | |
| # 'num': num, | |
| # 'dateRestrict': date_restrict, | |
| } | |
| request.update(params) | |
| return request | |
| def make_request(request): | |
| response = requests.get('https://www.googleapis.com/customsearch/v1', params=request) | |
| # print(f"Status Code: {response.status_code}") # Print status code for debugging | |
| if response.status_code != 200: | |
| raise Exception("Google Custom Search request failed with status code {}".format(response.status_code)) | |
| search_results = response.json() | |
| # print(f"Response JSON: {search_results}") # Print raw response for debugging | |
| if 'items' in search_results and len(search_results['items']) > 0: | |
| first_result_url = search_results['items'][0]['link'] | |
| return first_result_url | |
| else: | |
| return None | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| selected_shop = gr.Dropdown(shops.keys(), label="Shop") | |
| # with gr.Row(): | |
| # input_type = gr.Radio(["PDF", "Image"], label="Input Type", value="PDF") | |
| confidence_threshold = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, label="Confidence Threshold", value=0.5) | |
| with gr.Row(): | |
| pdf_input = gr.File(label="Input page (PDF)", file_types=["pdf"]) | |
| img_output = gr.AnnotatedImage(label="Annotated page (image)") | |
| section_btn = gr.Button("Identify Hotspots") | |
| # selected_section = gr.Textbox(label="Selected Hotspot") | |
| selected_section = gr.Markdown(label="Selected Hotspot") | |
| def get_yolo_predictions(image, confidence_threshold): | |
| # Use Roboflow API to get predictions | |
| predictions = model.predict(image, confidence=confidence_threshold, overlap=10).json() | |
| # print("predictions", predictions) | |
| return predictions | |
| def get_yolo_predictions_local(image, confidence_threshold): | |
| # Perform object detection on an image using the model | |
| predictions = model(image, conf=confidence_threshold) | |
| # print("predictions", predictions) | |
| return predictions | |
| def ocr(sub_image): | |
| # Use OCR to get text from image | |
| grey = cv2.cvtColor(sub_image, cv2.COLOR_BGR2GRAY) | |
| text = pytesseract.image_to_string(grey, lang='eng') | |
| return text | |
| def section(pdf, selected_shop, confidence_threshold): | |
| hostpots = [] | |
| # texts = [] | |
| sku = [] | |
| image, width_pdf, height_pdf, first_page = pdf_to_image(pdf) | |
| image_name ="test.jpg" | |
| image.save(image_name) | |
| # # Load a text page helper | |
| textpage = first_page.get_textpage() | |
| text_all = textpage.get_text_range() | |
| print("text_all", text_all) | |
| #print("selected_shop:", selected_shop) | |
| search_url = shops[selected_shop]['url'] | |
| ref_number_regex = shops[selected_shop]['regex'] | |
| # print("search_url:", search_url) | |
| brands = shops[selected_shop]['brands'] if 'brands' in shops[selected_shop] else None | |
| if brands: | |
| brands = [brand.upper() for brand in brands] | |
| if mode == 'rbf': | |
| predictions = get_yolo_predictions(image) | |
| else: | |
| # local inference | |
| results = get_yolo_predictions_local(image, confidence_threshold) | |
| for result in results: | |
| coordinates = result.boxes.xyxyn | |
| coordinates = coordinates.to("cpu") | |
| coordinates = coordinates.numpy() | |
| print("FOUND {} objects".format(len(coordinates))) | |
| # print("coordinates", coordinates) | |
| for idx, coord in enumerate(coordinates): | |
| # crop the image to the bounding box | |
| # expansion_factor = 0.01 | |
| x = coord[0] | |
| y = coord[1] | |
| w = coord[2] - coord[0] | |
| h = coord[3] - coord[1] | |
| # sub_image = image[y:y + h, x:x + w] | |
| print("x, y, w, h", x, y, w, h) | |
| print('width_pdf', width_pdf) | |
| print('height_pdf', height_pdf) | |
| enlarge_factor = 0.01 | |
| pdf_x = x * width_pdf | |
| pdf_y = y * height_pdf + enlarge_factor * height_pdf | |
| pdf_w = w * width_pdf + enlarge_factor * width_pdf | |
| pdf_h = h * height_pdf | |
| print("PDF", pdf_x, pdf_y, pdf_w, pdf_h) | |
| img_x = int(x * image.width) | |
| img_y = int(y * image.height) | |
| img_w = int(w * image.width) | |
| img_h = int(h * image.height) | |
| print("IMG", img_x, img_y, img_w, img_h) | |
| if shops[selected_shop]['url'] is None or shops[selected_shop]['url'] == "": | |
| # just append the hotspot to hostpots | |
| hostpots.append(((img_x, img_y, img_x + img_w, img_y + img_h), str(idx))) | |
| continue | |
| print('---------------------') | |
| # Extract text | |
| # text = textpage.get_text_bounded(left=pdf_x - enlarge_factor * width_pdf, | |
| # bottom=height_pdf - pdf_y - pdf_h - (enlarge_factor*height_pdf), | |
| # right=pdf_x+pdf_w + (enlarge_factor*width_pdf), | |
| # top=height_pdf-pdf_y + (enlarge_factor*height_pdf)) | |
| text = textpage.get_text_bounded(left=pdf_x + 50, | |
| bottom=height_pdf - pdf_y - pdf_h, | |
| right=pdf_x + pdf_w + 50, | |
| top=height_pdf - pdf_y) | |
| print("TEXT", text) | |
| # for object_number, text in enumerate(texts, start=1): | |
| #print("object", object_number) | |
| #print("text", text) | |
| ref_numbers = [] | |
| if selected_shop == "Decathlon": | |
| full_match = re.search(ref_number_regex, text) | |
| if full_match: | |
| #print("FULL MATCH: ", full_match) | |
| # Extracting all numbers from the captured sequence | |
| numbers_string = full_match.group(1) | |
| ref_numbers = re.findall(r'\d+', numbers_string) | |
| else: | |
| ref_numbers = re.findall(ref_number_regex, text) | |
| if selected_shop == "Colruyt": | |
| ref_numbers = [' '] | |
| print("REF NUMBERS: ", ref_numbers) | |
| if selected_shop == "Canadian Tire": | |
| temp_ref_numbers = ref_numbers | |
| found_skus_new = [] | |
| ref_numbers = [] | |
| split = False | |
| for ref_sku in temp_ref_numbers: | |
| if '/' in ref_sku: | |
| split = True | |
| sku_split = ref_sku.split('/') | |
| sku_last = sku_split[-1] | |
| length = len(sku_last) | |
| base_number = sku_split[0][:-length] | |
| ref_numbers.append(sku_split[0]) | |
| # base number + sku for each sku in sku_split except the first one | |
| for i in range(1, len(sku_split)): | |
| ref_numbers.append(base_number + sku_split[i]) | |
| if split == False: | |
| ref_numbers = temp_ref_numbers | |
| #print("REF NUMBERS: ", ref_numbers) | |
| matches = re.findall(regex_price, text) | |
| prices = [] | |
| for match in matches: | |
| if match[0]: # If the first group is not empty, the format is "price currency" | |
| prices.append(match[0] + match[1]) | |
| else: # Otherwise, the format is "currency price" | |
| prices.append(match[2] + match[3]) | |
| # price = prices[0] if len(prices) > 0 else "" | |
| #print("PRICES: ", prices) | |
| if brands: | |
| # print("BRANDS: yes") | |
| # Extracting the product name | |
| lines = text.split('\n') | |
| product_name = "" | |
| found_brand = False | |
| name_started = False | |
| for line in lines: | |
| if any(brand in line for brand in brands): | |
| found_brand = True | |
| #print("FOUND BRAND", line) | |
| continue | |
| if found_brand: | |
| #print('current line:', line) | |
| if line.strip() == "": # Skip blank lines | |
| continue | |
| if name_started and (line[0].isupper() or any(char.isdigit() for char in line)): | |
| break | |
| product_name += line.strip() + " " | |
| name_started = True | |
| product_name = product_name.strip() | |
| else: | |
| # Regex to match sequences of uppercase letters | |
| uppercase_words = re.findall(r'\b[A-Z]+\b', text) | |
| # Joining the extracted words with a space | |
| product_name = ' '.join(uppercase_words) | |
| #print("PRODUCT NAME: ", product_name) | |
| print("FOUND {} ref numbers".format(len(ref_numbers))) | |
| product_name | |
| for ref in ref_numbers: | |
| if search_google: | |
| # Use Google Search API to find items with the SKU | |
| to_seach = str(selected_shop) + ' ' + str(ref) | |
| if selected_shop == "Colruyt": | |
| # keep only the first 2 lines of text to search | |
| lines = text.split('\n') | |
| text = lines[0] + ' ' + lines[1] | |
| product_name = lines[0] | |
| to_seach = str(selected_shop) + ' ' + text | |
| print("to_search: ", to_seach) | |
| request = build_request(to_seach) | |
| search_query = make_request(request) | |
| print(f"First Result URL: {search_query}") | |
| if search_query is None: | |
| search_query = search_url + urllib.parse.quote(ref) | |
| else: | |
| search_query = search_url + urllib.parse.quote(ref) | |
| if ref not in sku: | |
| sku.append(ref) | |
| # if search_query not in hotspots_sku_links: | |
| hotspots_sku_links.append([{'sku': ref, 'url': search_query, 'prices': prices, 'product_name': product_name}]) | |
| hostpots.append(((img_x, img_y, (img_x + img_w), (img_h+img_y)), str(ref))) | |
| print("----- ADDED ------") | |
| if selected_shop == "Colruyt": | |
| sku.append(product_name) | |
| # if search_query not in hotspots_sku_links: | |
| hotspots_sku_links.append( | |
| [{'sku': product_name, 'url': search_query, 'prices': prices, 'product_name': product_name}]) | |
| hostpots.append(((img_x, img_y, (img_x + img_w), (img_h + img_y)), str(ref))) | |
| print("----- ADDED ------") | |
| if len(ref_numbers) == 0: | |
| hostpots.append(((img_x, img_y, (img_x + img_w), (img_h+img_y)), str(idx))) | |
| hotspots_sku_links.append([{'sku': None, 'url': None, 'prices': None, 'product_name': None}]) | |
| print("++++++ ADDED +++++") | |
| print("hotspots_sku_links: ", hotspots_sku_links) | |
| print("skus: ", sku) | |
| return (image, hostpots) | |
| section_btn.click(section, [pdf_input, selected_shop, confidence_threshold], img_output) | |
| def select_section(evt: gr.SelectData): | |
| hotspot_data = hotspots_sku_links[evt.index] | |
| markdown_output = "" | |
| for data in hotspot_data: | |
| url = data['url'] | |
| price_str = ", ".join(data['prices']) if data['prices'] else "No price found" | |
| product_name = data['product_name'] if data['product_name'] else "No product name found" | |
| if product_name != "No product name found": | |
| markdown_output += f"[{product_name}]({url})\n\n" | |
| else: | |
| markdown_output += f"[SKU: {data['sku']}]({url})\n\n" | |
| return markdown_output | |
| img_output.select(select_section, None, selected_section) | |
| if __name__ == "__main__": | |
| demo.launch(share=False) | |