Blame for app_file.py

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)