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 = [
    # the catalog's own products first: drinks, dairy, rice crackers
    "野菜生活100",
    "伊藤園 1日分の野菜",
    "Japanese vegetable juice shelf",
    "Japanese milk cartons supermarket",
    "meiji oishii milk",
    "せんべい 棚",
    "Japanese rice cracker shop display",
    "Japanese confectionery shelf",
    "Japanese supermarket beverage aisle",
    "Japanese convenience store refrigerator drinks",
    "Japanese supermarket dairy case",
    "スーパー 飲料 陳列",
    "お菓子 売り場",
    "Japanese snack aisle supermarket",
    "Japanese supermarket shelves drinks bottles",
    "Japanese supermarket interior shelves",
]

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:
    default = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                           "data", "external-shelves", "wikimedia-commons")
    root = sys.argv[1] if len(sys.argv) > 1 else default
    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, "not_a_photo": 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
        # A book scan is not a shelf. The search terms pull in .djvu volumes
        # whose page 1 is a title page, and Commons serves those as JPEG
        # thumbnails, so the extension check has to be on the source title.
        if any(title.lower().endswith(ext) for ext in (".djvu", ".pdf", ".svg", ".tif", ".tiff")):
            skipped["not_a_photo"] += 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):
        # Truncate the stem, never the extension: a 60-character cut on the
        # whole basename drops it, and the file then has no type for cv2 or
        # for anyone reading the directory.
        stem, extension = os.path.splitext(
            os.path.basename(urllib.parse.urlparse(entry["url"]).path).replace(" ", "_")
        )
        name = f"{index:03d}_{urllib.parse.unquote(stem)[:60]}{extension.lower()}"
        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()