Cv2 3833 compute and save sscd embeddings for images

.env_file

@@ -3,5 +3,7 @@ PRESTO_PORT=8000
 DEPLOY_ENV=local
 # MODEL_NAME=mean_tokens.Model
 MODEL_NAME=audio.Model
+# MODEL_NAME=image_sscd.Model
+# MODEL_NAME=image_pdq.Model
 AWS_ACCESS_KEY_ID=SOMETHING
 AWS_SECRET_ACCESS_KEY=OTHERTHING

.gitignore

@@ -1,2 +1,3 @@
 *.cpython-39.pyc
 *.pyc
+sscd_disc_mixup.torchscript.pt

Dockerfile

@@ -30,5 +30,6 @@ RUN pip install transformers
 RUN pip install pact-python
 RUN pip install --no-cache-dir -r requirements.txt
 RUN cd threatexchange/pdq/python && pip install .
+
 COPY . .
-CMD ["make", "run"]
+CMD ["make", "run"]

lib/model/generic_image.py

@@ -0,0 +1,27 @@
+from lib.model.model import Model
+
+from lib import schemas
+import urllib.request
+import io
+
+class GenericImageModel(Model):
+
+    def get_iobytes_for_image(self, image: schemas.Message) -> io.BytesIO:
+            """
+            Read file as bytes after requesting based on URL.
+            """
+            return io.BytesIO(
+                urllib.request.urlopen(
+                    urllib.request.Request(
+                        image.body.url,
+                        headers={'User-Agent': 'Mozilla/5.0'}
+                    )
+                ).read()
+            )
+
+    def process(self, image: schemas.Message) -> schemas.GenericItem:
+        """
+        Generic function for returning the actual response.
+        """
+
+        return self.compute_imagehash(self.get_iobytes_for_image(image))

lib/model/image_pdq.py

@@ -0,0 +1,20 @@
+from typing import Dict
+import io
+
+from lib.model.generic_image import GenericImageModel
+
+from pdqhashing.hasher.pdq_hasher import PDQHasher
+from lib import schemas
+
+class Model(GenericImageModel):
+    def compute_pdq(self, iobytes: io.BytesIO) -> str:
+        """Compute perceptual hash using ImageHash library
+        :param im: Numpy.ndarray
+        :returns: Imagehash.ImageHash
+        """
+        pdq_hasher = PDQHasher()
+        hash_and_qual = pdq_hasher.fromBufferedImage(iobytes)
+        return hash_and_qual.getHash().dumpBitsFlat()
+
+    def compute_imagehash(self, iobytes: io.BytesIO) -> str:
+        return self.compute_pdq(iobytes)

lib/model/image_sscd.py

@@ -0,0 +1,57 @@
+from typing import Dict
+import io
+
+from lib.model.generic_image import GenericImageModel
+from lib import schemas
+from torchvision import transforms
+import torch
+from lib.logger import logger
+import numpy as np
+from PIL import Image
+import urllib.request
+
+class Model(GenericImageModel):
+    def __init__(self):
+        super().__init__()
+        #FIXME: Load from a Meedan S3 bucket
+        try:
+            self.model = torch.jit.load("sscd_disc_mixup.torchscript.pt")
+        except:
+            logger.info("Downloading SSCD model...")
+            m=urllib.request.urlopen("https://dl.fbaipublicfiles.com/sscd-copy-detection/sscd_disc_mixup.torchscript.pt").read()
+            with open("sscd_disc_mixup.torchscript.pt","wb") as fh:
+                fh.write(m)
+            self.model = torch.jit.load("sscd_disc_mixup.torchscript.pt")
+        logger.info("SSCD model loaded")
+
+    def compute_sscd(self, iobytes: io.BytesIO) -> str:
+        """Compute perceptual hash using ImageHash library
+        :param im: Numpy.ndarray #FIXME
+        :returns: Imagehash.ImageHash #FIXME
+        """
+        # from SSCD-copy-detection readme https://github.com/facebookresearch/sscd-copy-detection/tree/main#preprocessing
+        # Normalization using the mean and std of Imagenet
+        normalize = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225],
+        )
+        # It is recommended by publishers of SSCD-copy-detection to preprocess images for inference either resizing the small edge to 288 or resizing the image to a square tensor.
+        # resizing the image to a square tensor is more effecient on gpus but can lead to skewed images and so loss of information. So, we are resizing the small edge to 288
+        small_288 = transforms.Compose([
+            transforms.Resize(288),
+            transforms.ToTensor(),
+            normalize,
+        ])
+        # Keeping the code example of resizing the image to a square tensor
+        # skew_320 = transforms.Compose([
+        #     transforms.Resize([320, 320]),
+        #     transforms.ToTensor(),
+        #     normalize,
+        # ])
+
+        image = Image.open(iobytes)
+        batch = small_288(image).unsqueeze(0)
+        embedding = self.model(batch)[0, :]
+        return np.asarray(embedding.detach().numpy()).tolist()
+
+    def compute_imagehash(self, iobytes: io.BytesIO) -> str:
+        return self.compute_sscd(iobytes)

requirements.txt

@@ -12,4 +12,8 @@ fasttext==0.9.2
 langcodes==3.3.0 
 requests==2.31.0
 pytest==7.4.0
-sentry-sdk==1.30.0
+sentry-sdk==1.30.0
+pytorch-lightning==1.5.10
+lightning-bolts==0.4.0
+torch==1.9.0
+torchvision==0.10.0

test/lib/model/test_image.py → test/lib/model/test_image_pdq.py

@@ -4,7 +4,7 @@
 from urllib.error import URLError
 from typing import Dict
 
-from lib.model.image import Model
+from lib.model.image_pdq import Model
 from lib import schemas
 
 class TestModel(unittest.TestCase):