[py-tx] Add a new match command line option (--rotations)

python-threatexchange/threatexchange/cli/match_cmd.py

@@ -9,18 +9,23 @@
 import logging
 import pathlib
 import typing as t
-
+import tempfile
 
 from threatexchange import common
 from threatexchange.cli.fetch_cmd import FetchCommand
 from threatexchange.cli.helpers import FlexFilesInputAction
 from threatexchange.exchanges.fetch_state import FetchedSignalMetadata
 
-from threatexchange.signal_type.index import IndexMatch, SignalTypeIndex
+from threatexchange.signal_type.index import (
+    IndexMatch,
+    SignalTypeIndex,
+    IndexMatchWithRotation,
+)
 from threatexchange.cli.exceptions import CommandError
 from threatexchange.signal_type.signal_base import BytesHasher, SignalType
 from threatexchange.cli.cli_config import CLISettings
-from threatexchange.content_type.content_base import ContentType
+from threatexchange.content_type.content_base import ContentType, RotationType
+from threatexchange.content_type.photo import PhotoContent
 
 from threatexchange.signal_type.signal_base import MatchesStr, TextHasher, FileHasher
 from threatexchange.cli import command_base
@@ -126,6 +131,12 @@ def init_argparse(cls, settings: CLISettings, ap: argparse.ArgumentParser) -> No
             action="store_true",
             help="show all matches, not just one per collaboration",
         )
+        ap.add_argument(
+            "--rotations",
+            "-R",
+            action="store_true",
+            help="for photos, generate and match all 8 simple rotations",
+        )
 
     def __init__(
         self,
@@ -136,6 +147,7 @@ def __init__(
         show_false_positives: bool,
         hide_disputed: bool,
         all: bool,
+        rotations: bool = False,
     ) -> None:
         self.content_type = content_type
         self.only_signal = only_signal
@@ -144,6 +156,7 @@ def __init__(
         self.hide_disputed = hide_disputed
         self.files = files
         self.all = all
+        self.rotations = rotations
 
         if only_signal and content_type not in only_signal.get_content_types():
             raise CommandError(
@@ -152,6 +165,11 @@ def __init__(
                 2,
             )
 
+        if self.rotations and not issubclass(content_type, PhotoContent):
+            raise CommandError(
+                "--rotations flag is only available for Photo content type", 2
+            )
+
     def execute(self, settings: CLISettings) -> None:
         if not settings.index.list():
             if not settings.in_demo_mode:
@@ -196,33 +214,86 @@ def execute(self, settings: CLISettings) -> None:
             for s_type, index in indices:
                 seen = set()  # TODO - maybe take the highest certainty?
                 if self.as_hashes:
-                    results = _match_hashes(path, s_type, index)
+                    results_from_hashes = _match_hashes(path, s_type, index)
+                    results: t.Sequence[IndexMatchWithRotation] = [
+                        IndexMatchWithRotation(match=match)
+                        for match in results_from_hashes
+                    ]
                 else:
-                    results = _match_file(path, s_type, index)
+                    results = _match_file(path, s_type, index, rotations=self.rotations)
 
                 for r in results:
-                    metadatas: t.List[t.Tuple[str, FetchedSignalMetadata]] = r.metadata
+                    if isinstance(r, IndexMatchWithRotation):
+                        metadatas: t.List[t.Tuple[str, FetchedSignalMetadata]] = (
+                            r.match.metadata
+                        )
+                        rotation_info = f" [{r.rotation_type.name}]"
+                        # Supposed to be without whitespace, but let's make sure
+                        distance_str = "".join(
+                            r.match.similarity_info.pretty_str().split()
+                        )
+
+                    elif isinstance(r, IndexMatch):
+                        metadatas = r.metadata
+                        rotation_info = ""
+                        distance_str = "".join(r.similarity_info.pretty_str().split())
+
                     for collab, fetched_data in metadatas:
                         if not self.all and collab in seen:
                             continue
                         seen.add(collab)
-                        # Supposed to be without whitespace, but let's make sure
-                        distance_str = "".join(r.similarity_info.pretty_str().split())
+
                         print(
                             s_type.get_name(),
-                            distance_str,
+                            distance_str + rotation_info,
                             f"({collab})",
                             fetched_data,
                         )
 
 
 def _match_file(
-    path: pathlib.Path, s_type: t.Type[SignalType], index: SignalTypeIndex
-) -> t.Sequence[IndexMatch]:
+    path: pathlib.Path,
+    s_type: t.Type[SignalType],
+    index: SignalTypeIndex,
+    rotations: bool = False,
+) -> t.Sequence[IndexMatchWithRotation]:
     if issubclass(s_type, MatchesStr):
-        return index.query(path.read_text())
+        matches = index.query(path.read_text())
+        return [IndexMatchWithRotation(match=match) for match in matches]
+
     assert issubclass(s_type, FileHasher)
-    return index.query(s_type.hash_from_file(path))
+
+    if not rotations or s_type != PhotoContent:
+        matches = index.query(s_type.hash_from_file(path))
+        return [IndexMatchWithRotation(match=match) for match in matches]
+
+    # Handle rotations for photos
+    with open(path, "rb") as f:
+        image_data = f.read()
+
+    rotated_images: t.Dict[RotationType, bytes] = PhotoContent.all_simple_rotations(
+        image_data
+    )
+    all_matches = []
+
+    for rotation_type, rotated_bytes in rotated_images.items():
+        # Create a temporary file to hold the image bytes
+        with tempfile.NamedTemporaryFile() as temp_file:
+            temp_file.write(rotated_bytes)
+            temp_file_path = pathlib.Path(temp_file.name)
+            matches = index.query(s_type.hash_from_file(temp_file_path))
+            temp_file_path.unlink()  # Clean up the temporary file
+
+        # Add rotation information if any matches were found
+        matches_with_rotations = []
+        for match in matches:
+            matches_with_rotations.append(
+                IndexMatchWithRotation(match=match, rotation_type=rotation_type)
+            )
+
+        all_matches.extend(matches_with_rotations)
+
+    return all_matches
 
 
 def _match_hashes(

python-threatexchange/threatexchange/content_type/content_base.py

@@ -7,7 +7,8 @@
 This records all the valid signal types for a piece of content.
 """
 
-from enum import Enum, auto
+
+from enum import Enum
 import typing as t
 
 from threatexchange import common
@@ -49,4 +50,4 @@ class RotationType(Enum):
     FLIPX = "flipx"  # Flip the object horizontally along the X-axis
     FLIPY = "flipy"  # Flip the object horizontally along the Y-axis
     FLIPPLUS1 = "flipplus1"  # Diagonal flip along the line y = x
-    FLIPMINUS1 = "flipminus1"  # Diagonal flip along the line y = -x
+    FLIPMINUS1 = "flipminus1"  # Diagonal flip along the line y = -x

python-threatexchange/threatexchange/content_type/photo.py

@@ -6,6 +6,7 @@
 """
 from PIL import Image
 import io
+import typing as t
 
 from .content_base import ContentType, RotationType
 
@@ -82,7 +83,7 @@ def flip_minus1(cls, image_data: bytes) -> bytes:
                 return buffer.getvalue()
 
     @classmethod
-    def all_simple_rotations(cls, image_data: bytes):
+    def all_simple_rotations(cls, image_data: bytes) -> t.Dict[RotationType, bytes]:
         """
         Generate the 8 naive rotations of an image.
 

python-threatexchange/threatexchange/signal_type/index.py

@@ -23,6 +23,7 @@
 import pickle
 import typing as t
 
+from threatexchange.content_type.content_base import RotationType
 
 T = t.TypeVar("T")
 S_Co = t.TypeVar("S_Co", covariant=True, bound="SignalSimilarityInfo")
@@ -139,6 +140,12 @@ def __eq__(self, other: t.Any) -> bool:
 IndexMatch = IndexMatchUntyped[SignalSimilarityInfo, T]
 
 
+@dataclass
+class IndexMatchWithRotation(t.Generic[T]):
+    match: IndexMatchUntyped[SignalSimilarityInfo, T]
+    rotation_type: RotationType = RotationType.ORIGINAL
+
+
 Self = t.TypeVar("Self", bound="SignalTypeIndex")