Implement image similarity searchHEADmain

4 files changed, 155 insertions, 13 deletions

diff --git a/README.md b/README.md
index e83d774..26fe25a 100644
--- a/README.md
+++ b/README.md
@@ -21,6 +21,8 @@ Options:
   -c, --classes NUMS    Nice class filter, comma-separated (e.g. 9,35)
   -s, --status STATUS   registered | pending | expired
   -v, --verbose         Show fetch timestamp column
+  -d, --download-logos  Download logo images to DIR
+      --similar-to IMG  Find trademarks with visually similar logos to IMG
   --json                Output raw JSON instead of table
   --page N              Page number (default: 1)
   --all-offices         Also include non-EU offices (US, GB, JP, CH, etc.)
@@ -35,6 +37,10 @@ tmview "APPLE" --classes 9,35 --status registered
 tmview "GOOGLE" --json | jq '.[:3]'
 tmview "COCA COLA" --limit 50 --page 2
 tmview "Hacking-Lab" --all-offices -v
+
+# Image similarity search
+tmview --similar-to logo.png
+tmview --similar-to logo.png --all-offices -d ./matches
 ```
 
 ## Output columns
diff --git a/src/tmview/cli.py b/src/tmview/cli.py
index 0b3948c..cebb95d 100644
--- a/src/tmview/cli.py
+++ b/src/tmview/cli.py
@@ -21,7 +21,7 @@ def main():
         prog="tmview",
         description="Search EU trademark registries via TMview",
     )
-    parser.add_argument("query", help="Trademark search term")
+    parser.add_argument("query", nargs="?", help="Trademark name to search for")
     parser.add_argument(
         "-o", "--offices",
         metavar="CODES",
@@ -50,6 +50,11 @@ def main():
         help="Show fetch timestamp column in table output",
     )
     parser.add_argument(
+        "--similar-to",
+        metavar="IMAGE",
+        help="Find trademarks with visually similar logos to IMAGE",
+    )
+    parser.add_argument(
         "-d", "--download-logos",
         metavar="DIR",
         help="Download trademark logo images to DIR",
@@ -80,6 +85,9 @@ def main():
 
     args = parser.parse_args()
 
+    if not args.query and not args.similar_to:
+        parser.error("provide a search query or --similar-to IMAGE")
+
     limit = min(args.limit, 100)
 
     if args.offices:
@@ -98,15 +106,25 @@ def main():
             sys.exit(1)
 
     try:
-        with console.status(f'[bold green]Searching TMview for "{args.query}"…[/bold green]'):
-            result = tm_client.search(
-                query=args.query,
-                offices=offices,
-                limit=limit,
-                page=args.page,
-                classes=classes,
-                status=args.status,
-            )
+        if args.similar_to:
+            status_msg = f"[bold green]Searching TMview for logos similar to {args.similar_to}…[/bold green]"
+            with console.status(status_msg):
+                result = tm_client.search_by_image(
+                    image_path=args.similar_to,
+                    offices=offices,
+                    limit=limit,
+                    page=args.page,
+                )
+        else:
+            with console.status(f'[bold green]Searching TMview for "{args.query}"…[/bold green]'):
+                result = tm_client.search(
+                    query=args.query,
+                    offices=offices,
+                    limit=limit,
+                    page=args.page,
+                    classes=classes,
+                    status=args.status,
+                )
     except RuntimeError as exc:
         console.print(f"[red]Error: {exc}[/red]")
         console.print("[dim]Suggestion: check your connection or retry in a few seconds.[/dim]")
@@ -118,7 +136,7 @@ def main():
     if args.json_output:
         tm_output.print_json(result)
     else:
-        tm_output.print_table(result, args.query, verbose=args.verbose)
+        tm_output.print_table(result, args.query, verbose=args.verbose, image_path=args.similar_to)
 
     if args.download_logos:
         trademarks = result["trademarks"]
diff --git a/src/tmview/client.py b/src/tmview/client.py
index 5e67802..6883848 100644
--- a/src/tmview/client.py
+++ b/src/tmview/client.py
@@ -1,4 +1,7 @@
+import hashlib
 import time
+from pathlib import Path
+
 import httpx
 from tmview import cache as _cache
 
@@ -50,6 +53,13 @@ _HEADERS = {
     "Origin": "https://www.tmdn.org",
 }
 
+_UPLOAD_HEADERS = {
+    "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/121.0 Safari/537.36",
+    "Referer": "https://www.tmdn.org/tmview/",
+    "Origin": "https://www.tmdn.org",
+    "Accept": "application/json",
+}
+
 
 def _build_payload(query, offices, limit, page, classes, status):
     payload = {
@@ -152,3 +162,107 @@ def search(query, offices=None, limit=20, page=1, classes=None, status=None):
             raise RuntimeError(
                 f"Connection failed after {attempt + 1} attempt(s): {exc}"
             ) from exc
+
+
+def _upload_image(image_path: str) -> dict:
+    path = Path(image_path)
+    with httpx.Client(timeout=60, headers=_UPLOAD_HEADERS) as client:
+        with path.open("rb") as fh:
+            resp = client.post(
+                f"{BASE_URL}/imageSearch/tm/uploadAndSegments",
+                files={"file": (path.name, fh, "image/jpeg")},
+                data={"clienttype": "desktop"},
+            )
+    try:
+        resp.raise_for_status()
+    except httpx.HTTPStatusError as exc:
+        raise RuntimeError(
+            f"Image upload failed (HTTP {exc.response.status_code}): {exc.response.text[:200]}"
+        ) from exc
+    data = resp.json()
+    seg = next((s for s in data.get("segments", []) if s.get("isSelected")), None)
+    if seg is None and data.get("segments"):
+        seg = data["segments"][0]
+    if seg is None:
+        raise RuntimeError("Image upload succeeded but returned no segments")
+    return {
+        "imageId": data["imageId"],
+        "imageName": data["imageName"],
+        "segmentLeft": str(seg["left"]),
+        "segmentRight": str(seg["right"]),
+        "segmentTop": str(seg["upper"]),
+        "segmentBottom": str(seg["lower"]),
+    }
+
+
+def search_by_image(image_path: str, offices=None, limit=20, page=1):
+    if offices is None:
+        offices = EU_OFFICES
+
+    image_bytes = Path(image_path).read_bytes()
+    image_hash = hashlib.sha256(image_bytes).hexdigest()
+    cache_params = {"image_sha256": image_hash, "offices": offices, "limit": limit, "page": page}
+    key = _cache.cache_key(cache_params)
+    cached = _cache.load(key)
+    if cached is not None:
+        result, fetched_at = cached
+        result["fetched_at"] = fetched_at
+        result["from_cache"] = True
+        return result
+
+    img = _upload_image(image_path)
+
+    payload = {
+        "imageId": img["imageId"],
+        "imageName": img["imageName"],
+        "segmentLeft": img["segmentLeft"],
+        "segmentRight": img["segmentRight"],
+        "segmentTop": img["segmentTop"],
+        "segmentBottom": img["segmentBottom"],
+        "colour": "false",
+        "criteria": "C",
+        "imageSearch": True,
+        "offices": offices,
+        "page": page,
+        "pageSize": limit,
+    }
+
+    delays = RETRY_DELAYS[:]
+    attempt = 0
+    while True:
+        try:
+            with httpx.Client(timeout=30) as client:
+                resp = client.post(
+                    f"{BASE_URL}/search/results",
+                    params={"translate": "true"},
+                    json=payload,
+                    headers=_HEADERS,
+                )
+            if resp.status_code in RETRYABLE_CODES and delays:
+                time.sleep(delays.pop(0))
+                attempt += 1
+                continue
+            resp.raise_for_status()
+            data = resp.json()
+            trademarks_raw = data.get("tradeMarks") or []
+            total = data.get("totalResults") or data.get("total") or len(trademarks_raw)
+            trademarks = [_parse_trademark(tm) for tm in trademarks_raw]
+            fetched_at = _cache.now_iso()
+            result = {"trademarks": trademarks, "total": total, "page": page}
+            _cache.save(key, result, fetched_at)
+            result["fetched_at"] = fetched_at
+            result["from_cache"] = False
+            return result
+        except httpx.HTTPStatusError as exc:
+            raise RuntimeError(
+                f"HTTP {exc.response.status_code} after {attempt + 1} attempt(s): "
+                f"{exc.response.text[:200]}"
+            ) from exc
+        except (httpx.TimeoutException, httpx.ConnectError, httpx.RemoteProtocolError) as exc:
+            if delays:
+                time.sleep(delays.pop(0))
+                attempt += 1
+                continue
+            raise RuntimeError(
+                f"Connection failed after {attempt + 1} attempt(s): {exc}"
+            ) from exc
diff --git a/src/tmview/output.py b/src/tmview/output.py
index 55bc3a4..95dc471 100644
--- a/src/tmview/output.py
+++ b/src/tmview/output.py
@@ -50,13 +50,17 @@ def _format_fetched(iso: str) -> str:
         return f"[dim]{iso}[/dim]"
 
 
-def print_table(result, query, verbose=False):
+def print_table(result, query, verbose=False, image_path=None):
     trademarks = result["trademarks"]
     total = result["total"]
     page = result["page"]
     fetched_at = result.get("fetched_at", "")
 
-    header = f'Trademark Search: "{query}" — {total:,} result{"s" if total != 1 else ""} (page {page})'
+    if image_path:
+        from pathlib import Path
+        header = f'Similar logos to "{Path(image_path).name}" — {total:,} result{"s" if total != 1 else ""} (page {page})'
+    else:
+        header = f'Trademark Search: "{query}" — {total:,} result{"s" if total != 1 else ""} (page {page})'
     console.print(f"\n[bold]{header}[/bold]\n")
 
     if not trademarks: