Feature/owlv2 support (#1)

README.md

@@ -117,6 +117,24 @@ NanoOWL runs real-time on Jetson Orin Nano.
 
 That's it!  If everything is working properly, you should see a visualization saved to ``data/owl_predict_out.jpg``.  
 
+<a id="Creating a Wheel Release and Tagging"></a>
+## Versioning and creating the Wheel
+After you have made all of the required edits and you want to create a release of the repository. 
+Edit the ``setup.py`` file to include the new version number that the wheel should be and commit these changes.
+The following command creates a wheel in the ``dist`` directory:
+```
+python setup.py bdist_wheel
+```
+
+Now that the wheel exists, make sure to tag the repository with the version number of the wheel.
+```
+git tag v<version-number> -a -m "A commit message to describe the changes and version"
+git tag
+git push origin v<version-number>
+```
+
+Now share the wheel one the server please.
+
 <a id="examples"></a>
 ## 🤸 Examples
 
@@ -213,6 +231,22 @@ live-edited text prompts.  To run the example
 
 
 
+### Example 4 - Fewshot prediction
+
+This example replicates the Image-Conditioned Detection example in the original OwlVit repo.  To run the example
+
+    ```bash
+    cd examples
+    python3 fewshot_predict.py \
+    --threshold="0.7,0.1" \
+    --image_encoder_engine=../data/owl_image_encoder_patch32.engine \
+    --query-image ../assets/cat_query_image.jpg ../assets/frog.jpg \
+    --query-label "a cat" "a frog"
+    ```
+
+By default the output will be saved to ``data/fewshot_predict_out.jpg``.
+
+
 <a id="acknowledgement"></a>
 ## 👏 Acknowledgement
 

examples/fewshot_predict.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import argparse
+import os.path
+import time
+
+import PIL.Image
+import torch
+from nanoowl.fewshot_predictor import FewshotPredictor
+from nanoowl.owl_drawing import draw_owl_output
+from nanoowl.owl_predictor import OwlPredictor
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image", type=str, default="../assets/cat_image.jpg")
+    parser.add_argument(
+        "--query-image",
+        metavar="N",
+        type=str,
+        nargs="+",
+        help="an example of what to look for in the image",
+        default=["../assets/frog.jpg", "../assets/cat_query_image.jpg"],
+    )
+    parser.add_argument(
+        "--query-label",
+        metavar="N",
+        type=str,
+        nargs="+",
+        help="a text label for each query image",
+        default=["a frog", "a cat"],
+    )
+    parser.add_argument("--threshold", type=str, default="0.1,0.7")
+    parser.add_argument("--output", type=str, default="../data/fewshot_predict_out.jpg")
+    parser.add_argument("--model", type=str, default="google/owlvit-base-patch32")
+    parser.add_argument(
+        "--image_encoder_engine",
+        type=str,
+        default="../data/owl_image_encoder_patch32.engine",
+    )
+    parser.add_argument("--profile", action="store_true")
+    parser.add_argument("--num_profiling_runs", type=int, default=30)
+    args = parser.parse_args()
+
+    image = PIL.Image.open(args.image)
+
+    query_images = []
+    for image_file in args.query_image:
+        if not os.path.isfile(image_file):
+            raise FileNotFoundError(f"File missing from {os.path.abspath(image_file)}")
+        else:
+            query_images.append(PIL.Image.open(image_file))
+
+    query_labels = args.query_label
+
+    thresholds = args.threshold.strip("][()")
+    thresholds = thresholds.split(",")
+    if len(thresholds) == 1:
+        thresholds = float(thresholds[0])
+    else:
+        thresholds = [float(x) for x in thresholds]
+
+    engine_path = (
+        args.image_encoder_engine if os.path.isfile(args.image_encoder_engine) else None
+    )
+    if not os.path.isfile(args.image_encoder_engine):
+        print(
+            f"No image encoder engine found at",
+            "{os.path.abspath(args.image_encoder_engine)}.",
+            "Continuing without tensorrt...",
+        )
+
+    predictor = FewshotPredictor(
+        owl_predictor=OwlPredictor(args.model, image_encoder_engine=engine_path)
+    )
+
+    query_embeddings = [
+        predictor.encode_query_image(image=query_image, text_hints=[query_labels])
+        for query_image, query_label in zip(query_images, query_labels)
+    ]
+
+    output = predictor.predict(image, query_embeddings, threshold=thresholds)
+
+    if args.profile:
+        torch.cuda.current_stream().synchronize()
+        t0 = time.perf_counter_ns()
+        for i in range(args.num_profiling_runs):
+            output = predictor.predict(image, query_embeddings, threshold=thresholds)
+        torch.cuda.current_stream().synchronize()
+        t1 = time.perf_counter_ns()
+        dt = (t1 - t0) / 1e9
+        print(f"PROFILING FPS: {args.num_profiling_runs/dt}")
+
+    image = draw_owl_output(image, output, text=query_labels, draw_text=True)
+
+    image.save(args.output)

nanoowl/build_image_encoder_engine.py

@@ -25,10 +25,12 @@
     parser.add_argument("--model_name", type=str, default="google/owlvit-base-patch32")
     parser.add_argument("--fp16_mode", type=bool, default=True)
     parser.add_argument("--onnx_opset", type=int, default=16)
+    parser.add_argument("--align_rois", type=bool, default=True)
     args = parser.parse_args()
 
     predictor = OwlPredictor(
-        model_name=args.model_name
+        model_name=args.model_name,
+        align_rois =args.align_rois,
     )
 
     predictor.build_image_encoder_engine(

nanoowl/fewshot_predictor.py

@@ -0,0 +1,168 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import List, Optional, Union
+
+import PIL.Image
+import torch
+
+from .image_preprocessor import ImagePreprocessor
+from .owl_predictor import (
+    OwlDecodeOutput,
+    OwlEncodeImageOutput,
+    OwlEncodeTextOutput,
+    OwlPredictor,
+)
+
+
+class FewshotPredictor(torch.nn.Module):
+    def __init__(
+        self,
+        owl_predictor: Optional[OwlPredictor] = None,
+        image_preprocessor: Optional[ImagePreprocessor] = None,
+        device: str = None,
+    ):
+        super().__init__()
+        device = device or "cuda" if torch.cuda.is_available() else "cpu"
+        self.owl_predictor = (
+            OwlPredictor(device=device) if owl_predictor is None else owl_predictor
+        )
+        self.image_preprocessor = (
+            ImagePreprocessor().to(device).eval()
+            if image_preprocessor is None
+            else image_preprocessor
+        )
+
+    @torch.no_grad()
+    def predict(
+        self,
+        image: PIL.Image,
+        query_embeddings: List,
+        threshold: Union[int, float, List[Union[int, float]]] = 0.1,
+        pad_square: bool = True,
+    ) -> OwlDecodeOutput:
+        image_tensor = self.image_preprocessor.preprocess_pil_image(image)
+
+        rois = torch.tensor(
+            [[0, 0, image.width, image.height]],
+            dtype=image_tensor.dtype,
+            device=image_tensor.device,
+        )
+
+        image_encodings = self.owl_predictor.encode_rois(
+            image_tensor, rois, pad_square=pad_square
+        )
+
+        return self.decode(image_encodings, query_embeddings, threshold)
+
+    def decode(
+        self,
+        image_output: OwlEncodeImageOutput,
+        query_embeds,
+        threshold: Union[int, float, List[Union[int, float]]] = 0.1,
+    ) -> OwlDecodeOutput:
+        num_input_images = image_output.image_class_embeds.shape[0]
+
+        image_class_embeds = image_output.image_class_embeds
+        image_class_embeds = image_class_embeds / (
+            torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6
+        )
+
+        if isinstance(threshold, (int, float)):
+            threshold = [threshold] * len(
+                query_embeds
+            )  # apply single threshold to all labels
+
+        query_embeds = torch.concat(query_embeds, dim=0)
+        query_embeds = query_embeds / (
+            torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6
+        )
+        logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)
+        logits = (logits + image_output.logit_shift) * image_output.logit_scale
+
+        scores_sigmoid = torch.sigmoid(logits)
+        scores_max = scores_sigmoid.max(dim=-1)
+        labels = scores_max.indices
+        scores = scores_max.values
+        masks = []
+        for i, thresh in enumerate(threshold):
+            label_mask = labels == i
+            score_mask = scores > thresh
+            obj_mask = torch.logical_and(label_mask, score_mask)
+            masks.append(obj_mask)
+        mask = masks[0]
+        for mask_t in masks[1:]:
+            mask = torch.logical_or(mask, mask_t)
+
+        input_indices = torch.arange(
+            0, num_input_images, dtype=labels.dtype, device=labels.device
+        )
+        input_indices = input_indices[:, None].repeat(1, self.owl_predictor.num_patches)
+
+        return OwlDecodeOutput(
+            labels=labels[mask],
+            scores=scores[mask],
+            boxes=image_output.pred_boxes[mask],
+            input_indices=input_indices[mask],
+        )
+
+    def encode_query_image(
+        self,
+        image: PIL.Image,
+        text_hints: List[str],
+        pad_square: bool = True,
+    ) -> torch.Tensor:
+        image_tensor = self.image_preprocessor.preprocess_pil_image(image)
+
+        text_encodings = self.encode_text(text_hints)
+
+        rois = torch.tensor(
+            [[0, 0, image.width, image.height]],
+            dtype=image_tensor.dtype,
+            device=image_tensor.device,
+        )
+
+        image_encodings = self.owl_predictor.encode_rois(
+            image_tensor, rois, pad_square=pad_square
+        )
+
+        return self.find_best_encoding(image_encodings, text_encodings)
+
+    def encode_text(self, texts: List[str]) -> OwlEncodeTextOutput:
+        return self.owl_predictor.encode_text(texts)
+
+    @staticmethod
+    def find_best_encoding(
+        image_output: OwlEncodeImageOutput,
+        text_output: OwlEncodeTextOutput,
+    ) -> torch.Tensor:
+        image_class_embeds = image_output.image_class_embeds
+        image_class_embeds = image_class_embeds / (
+            torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6
+        )
+        query_embeds = text_output.text_embeds
+        query_embeds = query_embeds / (
+            torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6
+        )
+        logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)
+        logits = (logits + image_output.logit_shift) * image_output.logit_scale
+
+        scores_sigmoid = torch.sigmoid(logits)
+        scores_max = scores_sigmoid.max(dim=-1)
+        scores = scores_max.values
+        best = torch.argmax(scores).item()
+        best_embed = image_class_embeds[:, best]
+        return best_embed

nanoowl/owl_drawing.py

@@ -14,13 +14,15 @@
 # limitations under the License.
 
 
-import PIL.Image
-import PIL.ImageDraw
+from typing import List
+
 import cv2
-from .owl_predictor import OwlDecodeOutput
 import matplotlib.pyplot as plt
 import numpy as np
-from typing import List
+import PIL.Image
+import PIL.ImageDraw
+
+from .owl_predictor import OwlDecodeOutput
 
 
 def get_colors(count: int):
@@ -36,7 +38,7 @@ def get_colors(count: int):
 def draw_owl_output(image, output: OwlDecodeOutput, text: List[str], draw_text=True):
     is_pil = not isinstance(image, np.ndarray)
     if is_pil:
-        image = np.asarray(image)
+        image = np.asarray(image).copy()
     font = cv2.FONT_HERSHEY_SIMPLEX
     font_scale = 0.75
     colors = get_colors(len(text))
@@ -48,13 +50,7 @@ def draw_owl_output(image, output: OwlDecodeOutput, text: List[str], draw_text=T
         box = [int(x) for x in box]
         pt0 = (box[0], box[1])
         pt1 = (box[2], box[3])
-        cv2.rectangle(
-            image,
-            pt0,
-            pt1,
-            colors[label_index],
-            4
-        )
+        cv2.rectangle(image, pt0, pt1, colors[label_index], 4)
         if draw_text:
             offset_y = 12
             offset_x = 0
@@ -66,9 +62,9 @@ def draw_owl_output(image, output: OwlDecodeOutput, text: List[str], draw_text=T
                 font,
                 font_scale,
                 colors[label_index],
-                2,# thickness
-                cv2.LINE_AA
+                2,  # thickness
+                cv2.LINE_AA,
             )
     if is_pil:
         image = PIL.Image.fromarray(image)
-    return image
+    return image