From 26fe24655c96e273c45db16b26050fd65d3c695a Mon Sep 17 00:00:00 2001 From: haianhng31 Date: Mon, 28 Oct 2024 21:01:18 +0000 Subject: [PATCH 1/5] Create RotationType enum --- .../threatexchange/content_type/content_base.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/python-threatexchange/threatexchange/content_type/content_base.py b/python-threatexchange/threatexchange/content_type/content_base.py index 7aade5342..880796991 100644 --- a/python-threatexchange/threatexchange/content_type/content_base.py +++ b/python-threatexchange/threatexchange/content_type/content_base.py @@ -7,6 +7,7 @@ This records all the valid signal types for a piece of content. """ +from enum import StrEnum, auto import typing as t from threatexchange import common @@ -32,3 +33,13 @@ def extract_additional_content( * Video => break out photo thumbnail, close caption text, audio """ return [] + +class RotationType(StrEnum): + ORIGINAL = auto() # No rotation; the object is in its original orientation + ROTATE90 = auto() # Rotates the object 90 degrees + ROTATE180 = auto() # Rotates the object 180 degrees (half-turn) + ROTATE270 = auto() # Rotates the object 270 degrees + FLIPX = auto() # Flip the object horizontally along the X-axis + FLIPY = auto() # Flip the object horizontally along the Y-axis + FLIPPLUS1 = auto() # Diagonal flip along the line y = x + FLIPMINUS1 = auto() # Diagonal flip along the line y = -x \ No newline at end of file From 931b81a49c1c1627cdf25f2c4addbb92391a11d6 Mon Sep 17 00:00:00 2001 From: haianhng31 Date: Mon, 28 Oct 2024 21:01:38 +0000 Subject: [PATCH 2/5] Create rotation helpers for PhotoContent --- .../threatexchange/content_type/photo.py | 76 ++++++++++++++++++- 1 file changed, 75 insertions(+), 1 deletion(-) diff --git a/python-threatexchange/threatexchange/content_type/photo.py b/python-threatexchange/threatexchange/content_type/photo.py index 0dcfcc379..6b63cbf87 100644 --- a/python-threatexchange/threatexchange/content_type/photo.py +++ b/python-threatexchange/threatexchange/content_type/photo.py @@ -4,9 +4,11 @@ """ Wrapper around the video content type. """ +from PIL import Image +import io from .content_base import ContentType - +from threatexchange.content_type.content_base import RotationType class PhotoContent(ContentType): """ @@ -19,3 +21,75 @@ class PhotoContent(ContentType): * frames from videos * thumbnails of videos """ + @classmethod + def rotate_image(cls, image_data: bytes, angle: float) -> bytes: + """ + Rotate an image by a given angle + """ + with Image.open(io.BytesIO(image_data)) as img: + rotated_img = img.rotate(angle, expand=True) + with io.BytesIO() as buffer: + rotated_img.save(buffer, format=img.format) + return buffer.getvalue() + + @classmethod + def flip_x(cls, image_data: bytes) -> bytes: + """ + Flip the image horizontally along the X-axis. + """ + with Image.open(io.BytesIO(image_data)) as img: + flipped_img = img.transpose(Image.FLIP_TOP_BOTTOM) + with io.BytesIO() as buffer: + flipped_img.save(buffer, format=img.format) + return buffer.getvalue() + + @classmethod + def flip_y(cls, image_data: bytes) -> bytes: + """ + Flip the image vertically along the Y-axis. + """ + with Image.open(io.BytesIO(image_data)) as img: + flipped_img = img.transpose(Image.FLIP_LEFT_RIGHT) + with io.BytesIO() as buffer: + flipped_img.save(buffer, format=img.format) + return buffer.getvalue() + + @classmethod + def flip_plus1(cls, image_data: bytes) -> bytes: + """ + Flip the image diagonally along the line y = x. + """ + with Image.open(io.BytesIO(image_data)) as img: + flipped_img = img.transpose(Image.Transpose.ROTATE_270).transpose(Image.FLIP_LEFT_RIGHT) + with io.BytesIO() as buffer: + flipped_img.save(buffer, format=img.format) + return buffer.getvalue() + + @classmethod + def flip_minus1(cls, image_data: bytes) -> bytes: + """ + Flip the image diagonally along the line y = -x. + """ + with Image.open(io.BytesIO(image_data)) as img: + flipped_img = img.transpose(Image.Transpose.ROTATE_90).transpose(Image.FLIP_LEFT_RIGHT) + with io.BytesIO() as buffer: + flipped_img.save(buffer, format=img.format) + return buffer.getvalue() + + @classmethod + def try_all_rotations(cls, image_data: bytes): + """ + Try all possible rotations to the image + """ + rotations = { + RotationType.ORIGINAL: image_data, + RotationType.ROTATE90: cls.rotate_image(image_data, 90), + RotationType.ROTATE180: cls.rotate_image(image_data, 180), + RotationType.ROTATE270: cls.rotate_image(image_data, 270), + RotationType.FLIPX: cls.flip_x(image_data), + RotationType.FLIPY: cls.flip_y(image_data), + RotationType.FLIPPLUS1: cls.flip_plus1(image_data), + RotationType.FLIPMINUS1: cls.flip_minus1(image_data) + } + return rotations + From dfe04d9a9df5f486a093efafb0eda04a48f7d753 Mon Sep 17 00:00:00 2001 From: haianhng31 Date: Wed, 30 Oct 2024 02:23:38 +0000 Subject: [PATCH 3/5] Include docstring and comments --- .../threatexchange/content_type/content_base.py | 5 +++++ .../threatexchange/content_type/photo.py | 10 +++++++--- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python-threatexchange/threatexchange/content_type/content_base.py b/python-threatexchange/threatexchange/content_type/content_base.py index 880796991..438e85137 100644 --- a/python-threatexchange/threatexchange/content_type/content_base.py +++ b/python-threatexchange/threatexchange/content_type/content_base.py @@ -35,6 +35,11 @@ def extract_additional_content( return [] class RotationType(StrEnum): + """ + Enum for 8 simple rotations of an image. + Used to store all generated rotations of an image, + whose algorithms don't have a native way to generate rotations during hashing. + """ ORIGINAL = auto() # No rotation; the object is in its original orientation ROTATE90 = auto() # Rotates the object 90 degrees ROTATE180 = auto() # Rotates the object 180 degrees (half-turn) diff --git a/python-threatexchange/threatexchange/content_type/photo.py b/python-threatexchange/threatexchange/content_type/photo.py index 6b63cbf87..988bb930f 100644 --- a/python-threatexchange/threatexchange/content_type/photo.py +++ b/python-threatexchange/threatexchange/content_type/photo.py @@ -24,7 +24,7 @@ class PhotoContent(ContentType): @classmethod def rotate_image(cls, image_data: bytes, angle: float) -> bytes: """ - Rotate an image by a given angle + Rotate an image by a given angle. """ with Image.open(io.BytesIO(image_data)) as img: rotated_img = img.rotate(angle, expand=True) @@ -77,9 +77,13 @@ def flip_minus1(cls, image_data: bytes) -> bytes: return buffer.getvalue() @classmethod - def try_all_rotations(cls, image_data: bytes): + def all_simple_rotations(cls, image_data: bytes): """ - Try all possible rotations to the image + Generate the 8 naive rotations of an image. + + This can be helpful for testing. + And for image algorithms that don't have a native way to generate rotations during hashing, + this can be a way to brute force rotations. """ rotations = { RotationType.ORIGINAL: image_data, From bd532aec83398363856fc2141aae45f0fc740bfb Mon Sep 17 00:00:00 2001 From: haianhng31 Date: Wed, 30 Oct 2024 17:19:34 +0000 Subject: [PATCH 4/5] Fix lint and small change to photo rotation helpers --- .../content_type/content_base.py | 20 ++++---- .../threatexchange/content_type/photo.py | 46 ++++++++++--------- 2 files changed, 36 insertions(+), 30 deletions(-) diff --git a/python-threatexchange/threatexchange/content_type/content_base.py b/python-threatexchange/threatexchange/content_type/content_base.py index 438e85137..bcd2b0c2f 100644 --- a/python-threatexchange/threatexchange/content_type/content_base.py +++ b/python-threatexchange/threatexchange/content_type/content_base.py @@ -34,17 +34,19 @@ def extract_additional_content( """ return [] + class RotationType(StrEnum): """ Enum for 8 simple rotations of an image. - Used to store all generated rotations of an image, + Used to store all generated rotations of an image, whose algorithms don't have a native way to generate rotations during hashing. """ - ORIGINAL = auto() # No rotation; the object is in its original orientation - ROTATE90 = auto() # Rotates the object 90 degrees - ROTATE180 = auto() # Rotates the object 180 degrees (half-turn) - ROTATE270 = auto() # Rotates the object 270 degrees - FLIPX = auto() # Flip the object horizontally along the X-axis - FLIPY = auto() # Flip the object horizontally along the Y-axis - FLIPPLUS1 = auto() # Diagonal flip along the line y = x - FLIPMINUS1 = auto() # Diagonal flip along the line y = -x \ No newline at end of file + + ORIGINAL = auto() # No rotation; the object is in its original orientation + ROTATE90 = auto() # Rotates the object 90 degrees + ROTATE180 = auto() # Rotates the object 180 degrees (half-turn) + ROTATE270 = auto() # Rotates the object 270 degrees + FLIPX = auto() # Flip the object horizontally along the X-axis + FLIPY = auto() # Flip the object horizontally along the Y-axis + FLIPPLUS1 = auto() # Diagonal flip along the line y = x + FLIPMINUS1 = auto() # Diagonal flip along the line y = -x diff --git a/python-threatexchange/threatexchange/content_type/photo.py b/python-threatexchange/threatexchange/content_type/photo.py index 988bb930f..6bdc50022 100644 --- a/python-threatexchange/threatexchange/content_type/photo.py +++ b/python-threatexchange/threatexchange/content_type/photo.py @@ -7,8 +7,8 @@ from PIL import Image import io -from .content_base import ContentType -from threatexchange.content_type.content_base import RotationType +from .content_base import ContentType, RotationType + class PhotoContent(ContentType): """ @@ -21,25 +21,26 @@ class PhotoContent(ContentType): * frames from videos * thumbnails of videos """ + @classmethod - def rotate_image(cls, image_data: bytes, angle: float) -> bytes: + def rotate_image(cls, image_data: bytes, angle: float) -> bytes: """ Rotate an image by a given angle. """ - with Image.open(io.BytesIO(image_data)) as img: + with Image.open(io.BytesIO(image_data)) as img: rotated_img = img.rotate(angle, expand=True) - with io.BytesIO() as buffer: + with io.BytesIO() as buffer: rotated_img.save(buffer, format=img.format) return buffer.getvalue() @classmethod - def flip_x(cls, image_data: bytes) -> bytes: + def flip_x(cls, image_data: bytes) -> bytes: """ Flip the image horizontally along the X-axis. """ - with Image.open(io.BytesIO(image_data)) as img: - flipped_img = img.transpose(Image.FLIP_TOP_BOTTOM) - with io.BytesIO() as buffer: + with Image.open(io.BytesIO(image_data)) as img: + flipped_img = img.transpose(Image.Transpose.FLIP_TOP_BOTTOM) + with io.BytesIO() as buffer: flipped_img.save(buffer, format=img.format) return buffer.getvalue() @@ -49,51 +50,54 @@ def flip_y(cls, image_data: bytes) -> bytes: Flip the image vertically along the Y-axis. """ with Image.open(io.BytesIO(image_data)) as img: - flipped_img = img.transpose(Image.FLIP_LEFT_RIGHT) + flipped_img = img.transpose(Image.Transpose.FLIP_LEFT_RIGHT) with io.BytesIO() as buffer: flipped_img.save(buffer, format=img.format) return buffer.getvalue() - @classmethod + @classmethod def flip_plus1(cls, image_data: bytes) -> bytes: """ Flip the image diagonally along the line y = x. """ with Image.open(io.BytesIO(image_data)) as img: - flipped_img = img.transpose(Image.Transpose.ROTATE_270).transpose(Image.FLIP_LEFT_RIGHT) + flipped_img = img.transpose(Image.Transpose.ROTATE_270).transpose( + Image.Transpose.FLIP_LEFT_RIGHT + ) with io.BytesIO() as buffer: flipped_img.save(buffer, format=img.format) return buffer.getvalue() - @classmethod + @classmethod def flip_minus1(cls, image_data: bytes) -> bytes: """ Flip the image diagonally along the line y = -x. """ with Image.open(io.BytesIO(image_data)) as img: - flipped_img = img.transpose(Image.Transpose.ROTATE_90).transpose(Image.FLIP_LEFT_RIGHT) + flipped_img = img.transpose(Image.Transpose.ROTATE_90).transpose( + Image.Transpose.FLIP_LEFT_RIGHT + ) with io.BytesIO() as buffer: flipped_img.save(buffer, format=img.format) return buffer.getvalue() - + @classmethod - def all_simple_rotations(cls, image_data: bytes): + def all_simple_rotations(cls, image_data: bytes): """ Generate the 8 naive rotations of an image. This can be helpful for testing. - And for image algorithms that don't have a native way to generate rotations during hashing, - this can be a way to brute force rotations. + And for image algorithms that don't have a native way to generate rotations during hashing, + this can be a way to brute force rotations. """ rotations = { - RotationType.ORIGINAL: image_data, + RotationType.ORIGINAL: image_data, RotationType.ROTATE90: cls.rotate_image(image_data, 90), RotationType.ROTATE180: cls.rotate_image(image_data, 180), RotationType.ROTATE270: cls.rotate_image(image_data, 270), RotationType.FLIPX: cls.flip_x(image_data), RotationType.FLIPY: cls.flip_y(image_data), RotationType.FLIPPLUS1: cls.flip_plus1(image_data), - RotationType.FLIPMINUS1: cls.flip_minus1(image_data) + RotationType.FLIPMINUS1: cls.flip_minus1(image_data), } return rotations - From 67bbcd22255da3a832e30983a18f6107f2ce1a5e Mon Sep 17 00:00:00 2001 From: haianhng31 Date: Thu, 31 Oct 2024 17:05:02 +0000 Subject: [PATCH 5/5] Change StrEnum to Enum --- .../content_type/content_base.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/python-threatexchange/threatexchange/content_type/content_base.py b/python-threatexchange/threatexchange/content_type/content_base.py index bcd2b0c2f..81e600877 100644 --- a/python-threatexchange/threatexchange/content_type/content_base.py +++ b/python-threatexchange/threatexchange/content_type/content_base.py @@ -7,7 +7,7 @@ This records all the valid signal types for a piece of content. """ -from enum import StrEnum, auto +from enum import Enum, auto import typing as t from threatexchange import common @@ -35,18 +35,18 @@ def extract_additional_content( return [] -class RotationType(StrEnum): +class RotationType(Enum): """ Enum for 8 simple rotations of an image. Used to store all generated rotations of an image, whose algorithms don't have a native way to generate rotations during hashing. """ - ORIGINAL = auto() # No rotation; the object is in its original orientation - ROTATE90 = auto() # Rotates the object 90 degrees - ROTATE180 = auto() # Rotates the object 180 degrees (half-turn) - ROTATE270 = auto() # Rotates the object 270 degrees - FLIPX = auto() # Flip the object horizontally along the X-axis - FLIPY = auto() # Flip the object horizontally along the Y-axis - FLIPPLUS1 = auto() # Diagonal flip along the line y = x - FLIPMINUS1 = auto() # Diagonal flip along the line y = -x + ORIGINAL = "original" # No rotation; the object is in its original orientation + ROTATE90 = "rotate90" # Rotates the object 90 degrees + ROTATE180 = "rotate180" # Rotates the object 180 degrees (half-turn) + ROTATE270 = "rotate270" # Rotates the object 270 degrees + 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