chore: add readme

README.md

@@ -3,3 +3,25 @@ license: apache-2.0
 ---
 
 refer: https://github.com/facebookresearch/sscd-copy-detection
+
+
+```python
+# code in src/demo.py
+import model
+from transformers import pipeline
+from transformers.image_utils import load_image
+
+pipe = pipeline(
+    task='sscd-copy-detection',
+    model='m3/sscd-copy-detection',
+    batch_size=10,
+    device='cpu',
+)
+
+vec1 = pipe(load_image("http://images.cocodataset.org/val2017/000000039769.jpg"))
+vec2 = pipe(load_image("http://images.cocodataset.org/val2017/000000039769.jpg"))
+
+import torch.nn.functional as F
+cos_sim = F.cosine_similarity(vec1, vec2, dim=0)
+print('similarity:', cos_sim.item())
+```

src/demo.py

@@ -0,0 +1,17 @@
+import model
+from transformers import pipeline
+from transformers.image_utils import load_image
+
+pipe = pipeline(
+    task='sscd-copy-detection',
+    model='m3/sscd-copy-detection',
+    batch_size=10,
+    device='cpu',
+)
+
+vec1 = pipe(load_image("http://images.cocodataset.org/val2017/000000039769.jpg"))
+vec2 = pipe(load_image("http://images.cocodataset.org/val2017/000000039769.jpg"))
+
+import torch.nn.functional as F
+cos_sim = F.cosine_similarity(vec1, vec2, dim=0)
+print('similarity:', round(cos_sim.item(), 3))

src/model.py

@@ -1,13 +1,15 @@
 from typing import List, Optional, Union
 from torchvision import transforms
 from PIL import Image
-
 from transformers.image_processing_utils import BaseImageProcessor
-from transformers import PreTrainedModel, PretrainedConfig
+from transformers import PretrainedConfig, PreTrainedModel, AutoConfig, AutoImageProcessor, AutoModel
 import os
 from huggingface_hub import hf_hub_download
 import torch
 import torch.nn as nn
+from transformers.pipelines import PIPELINE_REGISTRY
+from transformers.utils import add_end_docstrings
+from transformers.pipelines.base import Pipeline, build_pipeline_init_args
 class SscdImageProcessor(BaseImageProcessor):
     def __init__(
             self,
@@ -52,40 +54,69 @@ class SscdImageProcessor(BaseImageProcessor):
             image = image.convert('RGB')
         return preprocess(image).unsqueeze(0)
 
+
 class SscdConfig(PretrainedConfig):
     model_type = 'sscd-copy-detection'
+
     def __init__(self, model_path: str = None, **kwargs):
         if model_path is None:
             model_path = 'sscd_disc_mixup.torchscript.pt'
         super().__init__(model_path=model_path, **kwargs)
 
+
 class SscdModel(PreTrainedModel):
     config_class = SscdConfig
 
-    def __init__(self, config):
+    def __init__(self, config, model_path: str = None):
         super().__init__(config)
         self.dummy_param = nn.Parameter(torch.zeros(0))
-
-        print("______", config.name_or_path)
-
+        if model_path is None:
+            model_path = config.model_path
         is_local = os.path.isdir(config.name_or_path)
         if is_local:
             config.base_path = config.name_or_path
         else:
-            config_path = hf_hub_download(repo_id=config.name_or_path, filename='config.json')
-            config.base_path = os.path.dirname(config_path)
-        model_path =  config.base_path + '/' + config.model_path
-        print("___model_path___", model_path)
+            file_path = hf_hub_download(repo_id=config.name_or_path, filename=model_path)
+            config.base_path = os.path.dirname(file_path)
+        model_path = config.base_path + '/' + model_path
+        if model_path is not None:
+            self.model = torch.jit.load(model_path)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        return cls(AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs))
 
     def forward(self, inputs):
-        return self.model(inputs)
+        return self.model(inputs)[0, :]
+
+
+
+@add_end_docstrings(build_pipeline_init_args(has_image_processor=True))
+class SscdPipeline(Pipeline):
+    def __init__(self, model, **kwargs):
+        self.device_id = kwargs['device']
+        super().__init__(model=model, **kwargs)
 
-sscd_processor = SscdImageProcessor()
-sscd_processor.save_pretrained('new_model')
-sscd_config = SscdConfig(model_path='sscd_disc_mixup.torchscript.pt')
-sscd_config.save_pretrained('new_model')
+    def _sanitize_parameters(self, **kwargs):
+        return {}, {}, {}
+
+    def preprocess(self, input):
+        return self.image_processor.preprocess(input)
+
+    def _forward(self, inputs):
+        return self.model(inputs)
 
-model = SscdModel.from_pretrained('new_model')
+    def postprocess(self, model_outputs):
+        return model_outputs
 
 
+AutoConfig.register('sscd-copy-detection', SscdConfig)
+AutoModel.register(SscdConfig, SscdModel)
+AutoImageProcessor.register(SscdConfig, slow_image_processor_class=SscdImageProcessor)
+models = AutoModel.from_pretrained('m3/sscd-copy-detection')
 
+PIPELINE_REGISTRY.register_pipeline(
+    task='sscd-copy-detection',
+    pipeline_class=SscdPipeline,
+    pt_model=SscdModel
+)