Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[py-tx] Create rotation helper functions in PhotoContent using Pillow #1671

Merged
merged 5 commits into from
Nov 1, 2024
Merged

[py-tx] Create rotation helper functions in PhotoContent using Pillow #1671

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
26fe246
Create RotationType enum
haianhng31 Oct 28, 2024
931b81a
Create rotation helpers for PhotoContent
haianhng31 Oct 28, 2024
dfe04d9
Include docstring and comments
haianhng31 Oct 30, 2024
bd532ae
Fix lint and small change to photo rotation helpers
haianhng31 Oct 30, 2024
67bbcd2
Change StrEnum to Enum
haianhng31 Oct 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions python-threatexchange/threatexchange/content_type/content_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -32,3 +33,13 @@ def extract_additional_content(
* Video => break out photo thumbnail, close caption text, audio
"""
return []

class RotationType(StrEnum):
Dcallies marked this conversation as resolved.
Show resolved Hide resolved
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
76 changes: 75 additions & 1 deletion python-threatexchange/threatexchange/content_type/photo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
haianhng31 marked this conversation as resolved.
Show resolved Hide resolved

class PhotoContent(ContentType):
"""
Expand All @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Include a note on why we might want to do this :P

Namely, this is help for testing hash algorithms.

"""
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):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Let's go with all_simple_rotations

"""
Try all possible rotations to the image
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Something like

"""
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,
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

Loading