works with local yolo

app.py

@@ -9,12 +9,25 @@ import json
 import re
 from roboflow import Roboflow
 import gradio as gr
+from ultralytics import YOLO
 
-rf = Roboflow(api_key="TN586SQsdwjUUvMN3H8e")
-project = rf.workspace().project("smartflyers")
-model = project.version(9).model
 
-# Define shops with their respective search URL and reference number regex
+mode = 'local'
+# mode = 'rbf'
+
+if mode == 'local':
+    model = YOLO('/Users/selimgilon/Desktop/SportBoost/Projects/Iflyers/best-flyers-yolov8-1000img-v10.pt')
+else:
+    rf = Roboflow(api_key="TN586SQsdwjUUvMN3H8e")
+    project = rf.workspace().project("smartflyers")
+    model = project.version(9).model
+
+shop = None
+
+def process(selected_shop):
+    print(selected_shop)
+    return selected_shop
+
 shops = {
     'Latulippe': {
         'url': 'https://latulippe.com/fr/recherche/?keywords=',
@@ -24,365 +37,175 @@ shops = {
         'url': 'https://www.decathlon.fr/search?Ntt=',
         'regex': r'Réf. : (\d+)'
     },
-    'Metro' : {
-        'url': 'https://www.metro.ca/en/online-grocery/search?filter=',
-        'regex': r'(\d{6})'
-    },
-    'Walmart' : {
-        'url': 'https://www.walmart.ca/search?q=',
-        'regex': r'(\d{6})'
-    },
-    'Canadian Tire' : {
-        'url': 'https://www.canadiantire.ca/en/search-results.html?q=',
-        'regex': r'(\d{6})'
-    },
-    'Rona' : {
-        'url': 'https://www.rona.ca/en/search?searchType=product&q=',
-        'regex': r'(\d{6})'
-    },
-    'Home Depot' : {
-        'url': 'https://www.homedepot.ca/en/home/search.html?q=',
-        'regex': r'(\d{6})'
-    },
-    'BMR' : {
-        'url': 'https://www.bmr.co/fr/recherche.html?q=',
-        'regex': r'(\d{6})'
-    },
-    'Reno Depot' : {
-        'url': 'https://www.renodepot.com/en/search#?q=',
-        'regex': r'(\d{6})'
-    },
-    'Maxi' : {
-        'url': 'https://www.maxi.ca/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Super C' : {
-        'url': 'https://www.superc.ca/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Provigo' : {
-        'url': 'https://www.provigo.ca/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'IGA' : {
-        'url': 'https://www.iga.net/en/search?t={14A91276-40BA-49D5-940D-3692DE3DC381}&k=',
-        'regex': r'(\d{6})'
-    },
-    'Jean Coutu' : {
-        'url': 'https://www.jeancoutu.com/recherche?searchTerm=',
-        'regex': r'(\d{6})'
-    },
-    'Pharmaprix' : {
-        'url': 'https://www1.pharmaprix.ca/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Uniprix' : {
-        'url': 'https://www.uniprix.com/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Familiprix' : {
-        'url': 'https://www.familiprix.com/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'SAQ' : {
-        'url': 'https://www.saq.com/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Dollarama' : {
-        'url': 'https://www.dollarama.com/en-CA/search?text=',
-        'regex': r'(\d{6})'
-    },
-    'Canac' : {
-        'url': 'https://www.canac.ca/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Patrick Morin' : {
-        'url': 'https://www.patrickmorin.com/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Gagnon La Grande Quincaillerie' : {
-        'url': 'https://www.gagnonfreres.com/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Tanguay' : {
-        'url': 'https://www.tanguay.ca/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'The Source' : {
-        'url': 'https://www.thesource.ca/en-ca/search?text=',
-        'regex': r'(\d{6})'
-    },
-    'Best Buy' : {
-        'url': 'https://www.bestbuy.ca/en-ca/search?search=',
-        'regex': r'(\d{6})'
-    },
-    'Staples' : {
-        'url': 'https://www.staples.ca/search?query=',
-        'regex': r'(\d{6})'
-    },
-    'Bureau en Gros' : {
-        'url': 'https://www.staples.ca/search?query=',
-        'regex': r'(\d{6})'
-    },
-    'Centre Hi-Fi' : {
-        'url': 'https://www.centrehifi.com/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Stokes' : {
-        'url': 'https://www.stokesstores.com/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'RD Furniture' : {
-        'url': 'https://www.raymond.ca/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Go Sport' : {
-        'url': 'https://www.gosport.ca/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Sports Experts' : {
-        'url': 'https://www.sportsexperts.ca/en-CA/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'Atmosphere' : {
-        'url': 'https://www.atmosphere.ca/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'SAIL' : {
-        'url': 'https://www.sail.ca/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
-    'MEC' : {
-        'url': 'https://www.mec.ca/en/search?text=',
-        'regex': r'(\d{6})'
-    },
-    'Toys R Us' : {
-        'url': 'https://www.toysrus.ca/en/search?search-bar=',
-        'regex': r'(\d{6})'
-    },
     # Add more shops and their respective information here
 }
 
-def pdf_to_images(pdf_path):
-    images = convert_from_path(pdf_path)
-    return [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in images]
-
-
-
-def get_yolo_predictions(image):
-    # Convert image to suitable format if required
-    # Use Roboflow API to get predictions
-    predictions = model.predict(image, confidence=0.20, overlap=50).json()
-    return predictions
-
-def display_bbox_and_extract_text(image, predictions, page_number):
-    texts = []  # Store OCR text for each bbox
-    labeled_image = image.copy()  # Copy the image to draw labels on it
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    font_scale = 1  # Adjust the scale according to your image
-    color = (255, 255, 255)  # White color for text
-    bg_color = (0, 0, 0)  # Black background for text for visibility
-    thickness = 10 # Thickness of the bbox outline
-    text_thickness = 3  # Thickness of the font
-
-    for idx, pred in enumerate(predictions["predictions"]):
-        x_center = int(pred["x"])
-        y_center = int(pred["y"])
-        x = int(x_center - (pred["width"] / 2))
-        y = int(y_center - (pred["height"] / 2))
-        w = int(pred["width"])
-        h = int(pred["height"])
-
-        # Cropping the image to the bounding box
-        sub_image = image[y:y + h, x:x + w]
-
-        # Extract text using OCR
-        text = ocr(sub_image)
-        texts.append(text)
-
-        # Draw the bounding box on the image
-        start_point = (x, y)
-        end_point = (x + w, y + h)
-        bbox_color = (0, 255, 0)  # Green color in BGR for bbox
-        labeled_image = cv2.rectangle(labeled_image, start_point, end_point, bbox_color, thickness)
-
-        # Put the label number near the bounding box with background for visibility
-        label_text = str(idx + 1)
-        label_size = cv2.getTextSize(label_text, font, font_scale, text_thickness)[0]
-        label_position = (x, y + label_size[1] + 10) if y > label_size[1] else (x, y + label_size[1] + 20)
-
-        # Draw background rectangle for text
-        label_bg_start = (label_position[0], label_position[1] - label_size[1] - 10)
-        label_bg_end = (label_position[0] + label_size[0], label_position[1] + 10)
-        labeled_image = cv2.rectangle(labeled_image, label_bg_start, label_bg_end, bg_color, -1)
-
-        # Put the label number on the image
-        labeled_image = cv2.putText(labeled_image, label_text, label_position, font, font_scale, color, text_thickness)
-
-    # Draw the page number on the image
-    page_text = f"Page {page_number}"
-    page_text_size = cv2.getTextSize(page_text, font, font_scale, text_thickness)[0]
-    page_text_position = (image.shape[1] - page_text_size[0] - 20, page_text_size[1] + 20)  # top-right corner
-    col = (0, 0, 0)  # Black color in BGR for page number
-    labeled_image = cv2.putText(labeled_image, page_text, page_text_position, font, font_scale, col, text_thickness)
-
-    # Convert the labeled image from BGR to RGB for display
-    labeled_image = cv2.cvtColor(labeled_image, cv2.COLOR_BGR2RGB)
-    return labeled_image, texts
-
-
-def extract_reference_number(text):
-    # Look for "N°" followed by the reference number pattern
-    match = re.search(r'N°\s*([A-Za-z0-9]+)', text)
-    if match:
-        # Return the found reference number
-        return match.group(1)
-    else:
-        # Return None if no reference number is found
-        return None
-
-def ocr(sub_image):
-    # Use OCR to get text from image
-    text = pytesseract.image_to_string(sub_image, lang='eng')
-    return text
-
-
-# Set up your Google API key and Custom Search Engine ID
-GOOGLE_API_KEY = 'AIzaSyCD5eIu3N2uD6Vf5GE0VLBpkPsJ3mEZgA0'
-GOOGLE_CSE_ID = '32784ecc8a6bd4fad'
-
-
-def search_google(query):
-    search_url = "https://www.googleapis.com/customsearch/v1"
-    params = {
-        'q': query,
-        'key': GOOGLE_API_KEY,
-        'cx': GOOGLE_CSE_ID,
-        'num': 1  # Number of search results to return
-    }
-
-    response = requests.get(search_url, params=params)
-    search_results = response.json()
-    print(json.dumps(search_results, indent=2))
-
-    # Extract the first search result URL
-    url = search_results['items'][0]['link'] if 'items' in search_results else None
-    return url
-
-def search_latulippe(query):
-    # URLEncode the query to ensure it is HTTP-compliant.
-    encoded_query = urllib.parse.quote(query)
-    # Construct the search URL.
-    search_url = f"https://latulippe.com/fr/recherche/?keywords={encoded_query}"
-    return search_url
-
-
-def create_search_url(shop_url, query):
-    # URLEncode the query to ensure it is HTTP-compliant.
-    encoded_query = urllib.parse.quote(query)
-    # Construct the search URL.
-    if '{}' in shop_url:
-        # The user provided a format string with a placeholder for the query
-        search_url = shop_url.format(encoded_query)
-    else:
-        # Append the query to the base URL
-        search_url = f"{shop_url}{encoded_query}"
-    return search_url
-
-
-def clean_search_query(text):
-    # Add "Latulippe" to the beginning
-    text = "Latulippe " + text
-    # Keep only the first sentence by finding the first period followed by a space or the end of the text
-    first_sentence = re.split(r'(?<=[.!?]) +', text, maxsplit=1)[0]
-    # Convert to lowercase
-    first_sentence = first_sentence.lower()
-    # Replace newline characters with spaces
-    first_sentence = first_sentence.replace("\n", " ")
-    # Remove special characters except for basic punctuation
-    first_sentence = re.sub(r'[^a-z0-9\s,.]', '', first_sentence)
-    # Optionally, remove multiple consecutive spaces
-    first_sentence = re.sub(r'\s+', ' ', first_sentence).strip()
-    return first_sentence
-
-# Function to create a table from the extracted data
-def create_table(data):
-    # Assuming data is a list of tuples or lists with the following format:
-    # [(page_number, object_number, reference_code, url), (...), ...]
-    headers = ["Page", "Object", "Reference Code", "URL"]
-    # Convert URL to a clickable link
-    processed_data = []
-    for row in data:
-        clickable_link = f'<a href="{row[3]}" target="_blank">{row[3]}</a>'
-        processed_data.concat((row[0], row[1], row[2], clickable_link))
-    return gr.Dataframe(headers=headers, values=processed_data)
-
-def app(pdf_path, selected_shop):
-    if pdf_path is None:
-        raise gr.Error("No file uploaded. Please upload a PDF file.")
+hotspots_sku_links = []
+
+section_labels = [
+    "apple",
+    "banana",
+    "carrot",
+    "donut",
+    "eggplant",
+    "fish",
+    "grapes",
+    "hamburger",
+    "ice cream",
+    "juice",
+]
+
+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 not pdf_path.name.lower().endswith('.pdf'):
-        raise gr.Error("Invalid file format. Please upload a PDF file.")
 
     if selected_shop is None:
         raise gr.Error("Please select a shop.")
 
-    # Retrieve the selected shop's information
-    search_url = shops[selected_shop]['url']
-    ref_number_regex = shops[selected_shop]['regex']
-    table = pd.DataFrame(columns=["Page", "Object", "Reference", "URL"])
-    images = pdf_to_images(pdf_path)
-    output_images = []
-    all_urls = []
-
-    # Process the PDF and generate outputs if no errors are found
     try:
-        images = pdf_to_images(pdf_path)
-        for page_number, img in enumerate(images, start=1):
-            predictions = get_yolo_predictions(img)
-            img_with_bbox, texts = display_bbox_and_extract_text(img, predictions, page_number)
-            output_images.append(img_with_bbox)
-
-            # Extract reference numbers using regex and generate search URLs
-            #for text in texts:
-            for object_number, text in enumerate(texts, start=1):
-                ref_numbers = re.findall(ref_number_regex, text)
-                for ref in ref_numbers:
-                    search_query = search_url + urllib.parse.quote(ref)
-                    # all_urls.append({'page': page_number, 'object':object_number, 'reference': ref, 'url': search_query})
-                    new_row = pd.Series({
-                        "Page": page_number,
-                        "Object": object_number,
-                        "Reference": ref,
-                        "URL": search_query
-                    })
-                    table = pd.concat([table, pd.DataFrame([new_row])], ignore_index=True)
-
-        table_html = table.to_html(escape=False, index=False)
+        image_annotated = section(img)
         pass
     except Exception as e:
         return gr.Error(f"An error occurred: {str(e)}")
 
-    return output_images, table_html
-
+    return image_annotated, hotspots_sku_links[selected_shop]
 
 
-# Gradio interface definition
-iface = gr.Interface(
+iiface = gr.Interface(
     fn=app,
-    inputs=[
-        gr.File(label="Upload PDF"),  # Specify that you expect a file
-        gr.Dropdown(choices=list(shops.keys()), label="Select Shop")
-    ],
-    outputs=[
-        gr.Gallery(label="Processed Images"),  # Output for gallery
-        gr.HTML(label="Extracted Information")  # Output for table
-    ],
-    title="Automated Hotspotting",
-    description="Here is a demo of an automated hotspotting tool. Upload a PDF file and select the correponding shop to get started.",
-    #allow_flagging="never"
+    inputs=[gr.Dropdown(choices=list(shops.keys()), label="Select Shop"), gr.Image(label="Input page (image)", sources=["upload", "webcam"])],
+
+    outputs= [gr.AnnotatedImage(label="Annotated page (image)"), gr.Textbox(label="Selected Hotspot")],
+    title="Hotspot Identifier",
+    description="Identify hotspots on a page",
 )
 
-iface.launch(share=False)
 
+with gr.Blocks() as demo:
+
+
+    with gr.Row():
+        selected_shop = gr.Dropdown(shops.keys(), label="Shop")
+
+
+    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")
+
+
+    def get_yolo_predictions(image):
+        # Use Roboflow API to get predictions
+        predictions = model.predict(image, confidence=0.1, overlap=10).json()
+        # print("predictions", predictions)
+        return predictions
+
+
+    def get_yolo_predictions_local(image):
+        # Perform object detection on an image using the model
+        predictions = model(image, conf=0.1)
+        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):
+        hostpots = []
+        texts = []
+        sku = []
+
+        search_url = 'https://latulippe.com/fr/recherche/?keywords='
+        ref_number_regex = r'N° ?M(\d+)'
+
+        if mode == 'rbf':
+            predictions = get_yolo_predictions(img)
+            print("-------")
+            for idx, pred in enumerate(predictions["predictions"]):
+                x_center = int(pred["x"])
+                y_center = int(pred["y"])
+                x = int(x_center - (pred["width"] / 2))
+                y = int(y_center - (pred["height"] / 2))
+                w = int(pred["width"])
+                h = int(pred["height"])
+                # Cropping the image to the bounding box
+                sub_image = img[y:y + h, x:x + w]
+                cv2.imwrite(f"sub_image_{idx}.jpg", sub_image)
+
+                # Extract text using OCR
+                text = ocr(sub_image)
+                texts.append(text)
+
+                print("bbox", idx)
+                for object_number, text in enumerate(texts, start=1):
+                    print("object", object_number)
+                    print("text", text)
+                    ref_numbers = re.findall(ref_number_regex, text)
+                    for ref in ref_numbers:
+                        print("ref", ref)
+                        # append 'M' before the ref number
+                        ref = 'M' + ref
+                        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(search_query)
+                            hostpots.append(((x, y, x + w, y + h), ref))
+            return (img, hostpots)
+        else:
+            # get the height and width of the image
+
+            height, width, channels = img.shape
+            results = get_yolo_predictions_local(img)
+            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):
+                    print(coord.shape)
+                    print("coord", coord)
+                    # 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]
+
+                    # Extract text using OCR
+                    text = ocr(sub_image)
+                    texts.append(text)
+                    for object_number, text in enumerate(texts, start=1):
+                        print("object", object_number)
+                        print("text", text)
+                        ref_numbers = re.findall(ref_number_regex, text)
+                        for ref in ref_numbers:
+                            print("ref", ref)
+                            # append 'M' before the ref number
+                            ref = 'M' + ref
+                            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(search_query)
+                                    hostpots.append(((int(coord[0]), int(coord[1]), int(coord[2]), int(coord[3])), str(ref)))
+            return (img, hostpots)
+
+
+
+    section_btn.click(section, [img_input], img_output)
+
+    def select_section(evt: gr.SelectData):
+        return hotspots_sku_links[evt.index]
+
+    img_output.select(select_section, None, selected_section)
+
+if __name__ == "__main__":
+    demo.launch()