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,52 @@
+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
+        """
+        normalize = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225],
+        )
+        small_288 = transforms.Compose([
+            transforms.Resize(288),
+            transforms.ToTensor(),
+            normalize,
+        ])
+        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)

lib/queue/worker.py

@@ -51,10 +51,12 @@ def safely_respond(self, model: Model) -> List[schemas.Message]:
         responses = []
         if messages_with_queues:
             logger.debug(f"About to respond to: ({messages_with_queues})")
-            try:
-                responses = model.respond([schemas.Message(**{**json.loads(message.body), **{"model_name": model.model_name}}) for message, queue in messages_with_queues])
-            except Exception as e:
-                logger.error(e)
+            #try:
+            responses = model.respond([schemas.Message(**{**json.loads(message.body), **{"model_name": model.model_name}}) for message, queue in messages_with_queues])
+            logger.info("!!!!")
+            logger.info(responses)
+            #except Exception as e:
+            #    logger.error(e)
             self.delete_messages(messages_with_queues)
         return responses
 

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

@@ -15,7 +15,7 @@ def test_compute_pdq(self, mock_pdq_hasher):
             image_content = file.read()
         mock_hasher_instance = mock_pdq_hasher.return_value
         mock_hasher_instance.fromBufferedImage.return_value.getHash.return_value.dumpBitsFlat.return_value = '1001'
-        result = Model.compute_pdq(io.BytesIO(image_content))
+        result = Model().compute_pdq(io.BytesIO(image_content))
         self.assertEqual(result, '0011100000111011010110100001001110001011110100100010101011010111010110101010000111001010111000001010111111110000000101110010000011111110111110100100011111010010110110101111101100111001000000010010100101010111110001001101101011000110001000001110010000111100')
 
     @patch("urllib.request.urlopen")