product-recognition-demo/recognization/code/collect_shelf_images.py

"""
Collect openly licensed store-shelf photographs for evaluation.

Everything measured so far comes from two photographs. This pulls real store photos
from Wikimedia Commons - scriptable, licensed, and no authentication - so the
pipeline can be measured on images it has never seen.

Only permissively licensed files are kept (CC0, CC BY, CC BY-SA, public domain), and
each one is written to a manifest with its author, licence and source URL, so an
attribution page can be produced from it.

    python code/collect_shelf_images.py [output_dir]
"""

from __future__ import annotations

import json
import os
import ssl
import sys
import time
import urllib.parse
import urllib.request

import certifi

CONTEXT = ssl.create_default_context(cafile=certifi.where())
API = "https://commons.wikimedia.org/w/api.php"
USER_AGENT = "dorayaki-eval/0.1 (shelf recognition research; contact: repo owner)"
# Commons asks for a polite pace and offers thumbnails instead of originals. Both
# matter: a burst of original downloads got rate limited (HTTP 429) on the first run.
THUMBNAIL_WIDTH = 1600
REQUEST_DELAY = 1.2
MAX_RETRIES = 4

# Shelf photos of the categories the catalog covers: drinks and dairy, rice
# crackers and snacks, sauces and condiments.
QUERIES = [
    "supermarket shelf Japan",
    "Japanese supermarket beverages shelf",
    "Japanese convenience store drink shelf",
    "milk shelf Japan",
    "drink shelf Japan",
    "Japanese supermarket snack shelf",
    "rice crackers Japan shop",
    "せんべい 売り場",
    "飲料 棚 スーパー",
    "牛乳 棚",
    "スーパーマーケット 陳列",
    "Japanese grocery store interior",
    "Japanese supermarket dairy",
    "菓子 棚",
    "supermarket drinks Japan",
    "Japanese supermarket soft drinks",
]

PERMISSIVE = ("cc0", "cc-by", "cc by", "public domain", "pd-", "attribution")


def fetch(url: str, timeout: int = 60) -> bytes:
    """GET with backoff: Commons answers 429 when a client is too eager."""
    for attempt in range(MAX_RETRIES):
        try:
            request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
            with urllib.request.urlopen(request, timeout=timeout, context=CONTEXT) as response:
                return response.read()
        except urllib.error.HTTPError as error:
            if error.code != 429 or attempt == MAX_RETRIES - 1:
                raise
            time.sleep(REQUEST_DELAY * (2 ** attempt))
    raise RuntimeError("unreachable")


def api(params: dict) -> dict:
    url = API + "?" + urllib.parse.urlencode(params)
    time.sleep(REQUEST_DELAY)
    return json.loads(fetch(url, timeout=30).decode("utf-8"))


def search_files() -> dict:
    """Candidate files, with how many queries found each one."""
    found: dict = {}
    for query in QUERIES:
        try:
            result = api({"action": "query", "list": "search", "srsearch": query,
                          "srnamespace": "6", "srlimit": "25", "format": "json"})
        except Exception as error:
            print(f"  search failed for {query!r}: {error}")
            continue
        for hit in result.get("query", {}).get("search", []):
            title = hit["title"]
            if title.lower().endswith((".svg", ".pdf", ".ogv", ".webm")):
                continue
            found[title] = found.get(title, 0) + 1
        time.sleep(0.2)
    return found


def image_info(titles: list) -> dict:
    """URL, size and licence for each file, in batches of 50."""
    info: dict = {}
    for start in range(0, len(titles), 50):
        batch = titles[start:start + 50]
        try:
            result = api({"action": "query", "titles": "|".join(batch), "prop": "imageinfo",
                          "iiprop": "url|size|extmetadata", "iiurlwidth": str(THUMBNAIL_WIDTH),
                          "format": "json"})
        except Exception as error:
            print(f"  info failed for batch {start}: {error}")
            continue
        for page in (result.get("query", {}).get("pages") or {}).values():
            entries = page.get("imageinfo") or []
            if not entries:
                continue
            entry = entries[0]
            meta = entry.get("extmetadata") or {}
            info[page["title"]] = {
                "url": entry.get("thumburl") or entry.get("url"),
                "width": entry.get("width"),
                "height": entry.get("height"),
                "license": (meta.get("LicenseShortName") or {}).get("value", ""),
                "artist": (meta.get("Artist") or {}).get("value", ""),
                "description": (meta.get("ImageDescription") or {}).get("value", "")[:200],
                "page": f"https://commons.wikimedia.org/wiki/{urllib.parse.quote(page['title'].replace(' ', '_'))}",
            }
        time.sleep(0.2)
    return info


def main() -> None:
    root = sys.argv[1] if len(sys.argv) > 1 else os.path.expanduser("~/dorayaki-data/shelf-images")
    os.makedirs(root, exist_ok=True)

    print("searching Commons...")
    found = search_files()
    print(f"  {len(found)} distinct candidate files")

    info = image_info(sorted(found))
    print(f"  metadata for {len(info)} files")

    kept, skipped = [], {"license": 0, "size": 0, "missing": 0}
    for title, entry in sorted(info.items()):
        if not entry.get("url"):
            skipped["missing"] += 1
            continue
        licence = (entry.get("license") or "").lower()
        if not any(token in licence for token in PERMISSIVE):
            skipped["license"] += 1
            continue
        if (entry.get("width") or 0) < 1000 or (entry.get("height") or 0) < 700:
            skipped["size"] += 1
            continue
        kept.append((title, entry))

    print(f"  keeping {len(kept)} (skipped: {skipped})")
    manifest = []
    for index, (title, entry) in enumerate(kept):
        stem = os.path.basename(urllib.parse.urlparse(entry["url"]).path).replace(" ", "_")
        stem = urllib.parse.unquote(stem)[:60]
        name = f"{index:03d}_{stem}"
        path = os.path.join(root, name)
        if not os.path.exists(path):
            try:
                payload = fetch(entry["url"])
                with open(path, "wb") as handle:
                    handle.write(payload)
            except Exception as error:
                print(f"  download failed for {title}: {error}")
                continue
        manifest.append({
            "file": name,
            "title": title,
            "width": entry["width"],
            "height": entry["height"],
            "license": entry["license"],
            "artist": entry["artist"],
            "source": entry["page"],
            "queries": found.get(title, 0),
        })
        print(f"  [{index:>3}] {entry['width']}x{entry['height']} {entry['license']:<14} {title}")

    manifest_path = os.path.join(root, "manifest.json")
    with open(manifest_path, "w", encoding="utf-8") as handle:
        json.dump({"source": "Wikimedia Commons", "images": manifest}, handle, ensure_ascii=False, indent=1)
    print(f"\n{len(manifest)} images in {root}\nmanifest: {manifest_path}")


if __name__ == "__main__":
    main()