first commit

app.py

@@ -0,0 +1,148 @@
+import gzip
+import io
+import json
+import random
+import re
+import tempfile
+from typing import Dict, List, Optional
+
+from PIL import Image
+import requests
+import streamlit as st
+
+
+http_session = requests.Session()
+
+API_URL = "https://world.openfoodfacts.org/api/v0"
+PRODUCT_URL = API_URL + "/product"
+OFF_IMAGE_BASE_URL = "https://static.openfoodfacts.org/images/products"
+BARCODE_PATH_REGEX = re.compile(r"^(...)(...)(...)(.*)$")
+
+
+@st.cache()
+def load_nn_data(url: str):
+    r = http_session.get(url)
+    with gzip.open(io.BytesIO(r.content), "rt") as f:
+        content = f.read()
+        return {int(key): value for key, value in json.loads(content).items()}
+
+
+@st.cache()
+def load_logo_data(url: str):
+    r = http_session.get(url)
+    with gzip.open(io.BytesIO(r.content), "rt") as f:
+        content = f.read()
+        return {
+            int(item["id"]): item for item in (json.loads(x) for x in map(str.strip, content))
+        }
+
+
+def get_image_from_url(
+    image_url: str,
+    error_raise: bool = False,
+    session: Optional[requests.Session] = None,
+) -> Optional[Image.Image]:
+    if session:
+        r = http_session.get(image_url)
+    else:
+        r = requests.get(image_url)
+
+    if error_raise:
+        r.raise_for_status()
+
+    if r.status_code != 200:
+        return None
+
+    with tempfile.NamedTemporaryFile() as f:
+        f.write(r.content)
+        image = Image.open(f.name)
+
+    return image
+
+
+def split_barcode(barcode: str) -> List[str]:
+    if not barcode.isdigit():
+        raise ValueError("unknown barcode format: {}".format(barcode))
+
+    match = BARCODE_PATH_REGEX.fullmatch(barcode)
+
+    if match:
+        return [x for x in match.groups() if x]
+
+    return [barcode]
+
+
+def get_cropped_image(barcode: str, image_id: str, bounding_box):
+    image_path = generate_image_path(barcode, image_id)
+    url = OFF_IMAGE_BASE_URL + image_path
+    image = get_image_from_url(url, session=http_session)
+
+    if image is None:
+        return
+
+    ymin, xmin, ymax, xmax = bounding_box
+    (left, right, top, bottom) = (
+        xmin * image.width,
+        xmax * image.width,
+        ymin * image.height,
+        ymax * image.height,
+    )
+    return image.crop((left, top, right, bottom))
+
+
+def generate_image_path(barcode: str, image_id: str) -> str:
+    splitted_barcode = split_barcode(barcode)
+    return "/{}/{}.jpg".format("/".join(splitted_barcode), image_id)
+
+
+def display_predictions(
+    logo_data: Dict,
+    nn_data: Dict,
+    logo_id: Optional[int] = None,
+):
+    if not logo_id:
+        logo_id = random.choice(list(nn_data.keys()))
+
+    st.write(f"Logo ID: {logo_id}")
+    logo = logo_data[logo_id]
+    logo_nn_data = nn_data[logo_id]
+    nn_ids = logo_nn_data["ids"]
+    nn_distances = logo_nn_data["distances"]
+    annotation = logo_nn_data["annotation"]
+
+    cropped_image = get_cropped_image(
+        logo["barcode"], logo["image_id"], logo["bounding_box"]
+    )
+
+    if cropped_image is None:
+        return
+    st.image(cropped_image, annotation, width=200)
+
+    cropped_images: List[Image.Image] = []
+    captions: List[str] = []
+    for closest_id, distance in zip(nn_ids, nn_distances):
+        closest_logo = logo_data[closest_id]
+
+        cropped_image = get_cropped_image(
+            closest_logo["barcode"],
+            closest_logo["image_id"],
+            closest_logo["bounding_box"],
+        )
+        if cropped_image is None:
+            continue
+
+        if cropped_image.height > cropped_image.width:
+            cropped_image = cropped_image.rotate(90)
+
+        cropped_images.append(cropped_image)
+        captions.append(f"distance: {distance}")
+
+    if cropped_images:
+        st.image(cropped_images, captions, width=200)
+
+
+st.sidebar.title("Logo Nearest Neighbor Demo")
+# st.sidebar.write("")
+nn_data = load_nn_data("https://static.openfoodfacts.org/data/logos/exact_100_neighbours.json.gz")
+logo_data = load_logo_data("https://static.openfoodfacts.org/data/logos/logo_annotations.jsonl.gz")
+display_predictions(logo_data=logo_data, nn_data=nn_data)

requirements.txt

@@ -0,0 +1,3 @@
+Pillow==9.3.0
+requests==2.28.1
+streamlit==1.15.1