Spaces:

SelimGilon
/

Flyers

No application file

File size: 27,526 Bytes

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 = {
    '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:
        # find the first result with a 'linl' containing 'www.colruyt' and return it
        first_result_url = None
        for item in search_results['items']:
            if 'link' in item and 'www.colruyt' in item['link']:
                first_result_url = item['link']
                break
        #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")
    product_dataframe = gr.Dataframe(label="Products Found")

    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 = []
        product_table_data = []

        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):
                    print("==============")
                    # 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.0
                    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)

                    text = textpage.get_text_bounded(left=pdf_x,
                                                     bottom=height_pdf - pdf_y - pdf_h,
                                                     right=pdf_x + pdf_w,
                                                     top=height_pdf - pdf_y)


                    print("TEXT", text)


                    if search_google:
                        # Use Google Search API to find items with the SKU
                        # keep only the first 2 lines of text to search
                        lines = text.split('\n')
                        # delete the \r in all lines
                        lines = [line.replace('\r', '') for line in lines]
                        # remove the blank spaces at the beginning and end of each line
                        lines = [line.strip() for line in lines]
                        text = '\n'.join(lines[:2])
                        # if the text is less than 3 words, add the next line
                        if len(text.split(' ')) < 2:
                            text = '\n'.join(lines[:3])
                        product_name = text

                        to_seach =  str(selected_shop) + ' ' + text + str(selected_shop)
                        print("TO SEARCH: ", to_seach)
                        request = build_request(to_seach)
                        search_query = make_request(request)
                        print(f"First Result URL: {search_query}")
                        # search_query = "https://www.google.com/search?q="

                        sku.append(product_name)
                        # if search_query not in hotspots_sku_links:
                        hotspots_sku_links.append([{'sku': product_name, 'url': search_query, 'product_name': product_name}])
                        hostpots.append(((img_x, img_y, (img_x + img_w), (img_h+img_y)), str(product_name)))
                        print("----- ADDED ------")



            print("hotspots_sku_links: ", hotspots_sku_links)
            print("skus: ", sku)

            for data in hotspots_sku_links:
                for item in data:
                    product_table_data.append({
                        'Product Name': item['product_name'] if item['product_name'] else "Not Found",
                        'URL': item['url']
                    })


            # Create a DataFrame
            product_dataframe_data = pd.DataFrame(product_table_data)

            # flattern the list hotspots_sku_links
            # hotspots_sku_links = [item for sublist in hotspots_sku_links for item in sublist]
            print("hotspots_sku_links: ", hotspots_sku_links)
            print('length of hotspots_sku_links: ', len(hotspots_sku_links))
            assert len(hotspots_sku_links) == len(sku)
            return (image, hostpots), product_dataframe_data


    section_btn.click(section, [pdf_input, selected_shop, confidence_threshold], [img_output, product_dataframe])

    def select_section(evt: gr.SelectData):
        hotspot_data = hotspots_sku_links[evt.index]
        markdown_output = ""
        for data in hotspot_data:
            url = data['url']
            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)

demo.queue(max_size=99).launch(debug=True)

# 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
#
#
#
#
# with gr.Blocks() as demo:
#     search_google = False
#     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"
#     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 app(img, selected_shop):
#     #     if img is None:
#     #         raise gr.Error("No file uploaded. Please upload an image.")
#     #
#     #     # Check if the uploaded file is a PDF
#     #     if selected_shop is None:
#     #         raise gr.Error("Please select a shop.")
#     #
#     #     try:
#     #         image_annotated = section(img)
#     #         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.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():
#         img_input = gr.Image(label="Input page (image)", sources=["upload", "webcam"])
#         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")
#     product_dataframe = gr.Dataframe(label="Products Found")
#
#
#     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(img, selected_shop, confidence_threshold):
#         hostpots = []
#         sku = []
#         product_table_data = []
#
#         global hotspots_sku_links
#         hotspots_sku_links = []
#
#         #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(img)
#         else:
#             # local inference
#             results = get_yolo_predictions_local(img, confidence_threshold)
#             for result in results:
#                 coordinates = result.boxes.xyxy
#                 coordinates = coordinates.to("cpu")
#                 coordinates = coordinates.numpy()
#                 print("FOUND {} objects".format(len(coordinates)))
#                 for idx, coord in enumerate(coordinates):
#                     # crop the image to the bounding box
#                     x = int(coord[0])
#                     y = int(coord[1])
#                     w = int(coord[2] - coord[0])
#                     h = int(coord[3] - coord[1])
#                     sub_image = img[y:y + h, x:x + w]
#                     if shops[selected_shop]['url'] is None or shops[selected_shop]['url'] == "":
#                         print("No search url for this shop")
#                         # just append the hotspot to hostpots
#                         hostpots.append(((x, y, x + w, y + h), str(idx)))
#                         continue
#
#                     # Extract text using OCR
#                     text = ocr(sub_image)
#                     print("text: ", text)
#
#                     if selected_shop == "Latulippe":
#                         # if there is a 'i' between a "M" and a number, it is a sku. Replace the 'i' with a '1'
#                         text = re.sub(r'Mi(\d+)', r'M1\1', 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)
#                             print("REF NUMBERS: ", ref_numbers)
#                     else:
#                         ref_numbers = re.findall(ref_number_regex, text)
#                         # print("REF NUMBERS: ", ref_numbers)
#
#                     if selected_shop == "Canadian Tire":
#                         temp_ref_numbers = ref_numbers
#                         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
#
#                     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)
#
#                     for ref in ref_numbers:
#                         print("ref", ref)
#                         if search_google:
#                             # Use Google Search API to find items with the SKU
#                             to_seach = str(selected_shop) + ' ' + str(ref)
#                             # 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(((int(coord[0]), int(coord[1]), int(coord[2]), int(coord[3])), str(ref)))
#
#             for data in hotspots_sku_links:
#                 for item in data:
#                     product_table_data.append({
#                         'Reference Number': item['sku'],
#                         'Product Name': item['product_name'] if item['product_name'] else "Not Found",
#                         # 'Price': ", ".join(item['prices']) if item['prices'] else " ",
#                         'URL': item['url']
#                     })
#
#             # Create a DataFrame
#             product_dataframe_data = pd.DataFrame(product_table_data)
#
#             # flattern the list hotspots_sku_links
#             # hotspots_sku_links = [item for sublist in hotspots_sku_links for item in sublist]
#             print("hotspots_sku_links: ", hotspots_sku_links)
#             print('length of hotspots_sku_links: ', len(hotspots_sku_links))
#             print('shape of hotspots_sku_links: ', np.shape(hotspots_sku_links))
#             print("skus: ", sku)
#             print('length of skus: ', len(sku))
#             assert len(hotspots_sku_links) == len(sku)
#             return (img, hostpots), product_dataframe_data
#
#
#     section_btn.click(section, [img_input, selected_shop, confidence_threshold], [img_output, product_dataframe])
#
#
#
#     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":
#                 if price_str != "No price found":
#                     markdown_output += f"[{product_name}]({url}): {price_str}\n\n"
#                 else:
#                     markdown_output += f"[{product_name}]({url})\n\n"
#             else:
#                 if price_str != "No price found":
#                     markdown_output += f"[SKU: {data['sku']}]({url}): {price_str}\n\n"
#                 else:
#                     markdown_output += f"[SKU: {data['sku']}]({url})\n\n"
#         return markdown_output
#
#     img_output.select(select_section, None, selected_section)
#
# # demo.launch(share=False)
# demo.queue(max_size=99).launch(debug=True)