A project that optimizes OWL-ViT for real-time inference with NVIDIA TensorRT.
# 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 PIL.Image
import time
import torch
from nanoowl.owl_predictor import (
    OwlPredictor
)
from nanoowl.owl_drawing import (
    draw_owl_output
)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--image", type=str, default="../assets/owl_glove_small.jpg")
    parser.add_argument("--prompt", type=str, default="")
    parser.add_argument("--threshold", type=float, default=0.1)
    parser.add_argument("--output", type=str, default="../data/owl_predict_out.jpg")
    parser.add_argument("--model", type=str, default="google/owlvit-base-patch32")
    parser.add_argument("--image_encoder_engine", type=str, default="../data/owlvit_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()

    prompt = args.prompt.strip("][()")

    text = prompt.split(',')
    
    print(text)

    predictor = OwlPredictor(
        args.model,
        image_encoder_engine=args.image_encoder_engine
    )

    image = PIL.Image.open(args.image)
    
    text_encodings = predictor.encode_text(text)

    output = predictor.predict(
        image=image, 
        text=text, 
        text_encodings=text_encodings,
        threshold=args.threshold,
        pad_square=False
    )

    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=image, 
                text=text, 
                text_encodings=text_encodings,
                threshold=args.threshold,
                pad_square=False
            )
        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=text, draw_text=True)

    image.save(args.output)
examples/owl_predict.py
NVIDIA-AI-IOT/nanoowl
# 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 asyncio
import argparse
from aiohttp import web, WSCloseCode
import logging
import weakref
import cv2
import time
import PIL.Image
import matplotlib.pyplot as plt
from typing import List
from nanoowl.tree import Tree
from nanoowl.tree_predictor import (
    TreePredictor
)
from nanoowl.tree_drawing import draw_tree_output
from nanoowl.owl_predictor import OwlPredictor


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("image_encode_engine", type=str)
    parser.add_argument("--image_quality", type=int, default=50)
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--host", type=str, default="0.0.0.0")
    args = parser.parse_args()

    CAMERA_DEVICE = 0
    IMAGE_QUALITY = args.image_quality

    predictor = TreePredictor(
        owl_predictor=OwlPredictor(
            image_encoder_engine=args.image_encode_engine
        )
    )

    prompt_data = None

    def get_colors(count: int):
        cmap = plt.cm.get_cmap("rainbow", count)
        colors = []
        for i in range(count):
            color = cmap(i)
            color = [int(255 * value) for value in color]
            colors.append(tuple(color))
        return colors


    def cv2_to_pil(image):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return PIL.Image.fromarray(image)


    async def handle_index_get(request: web.Request):
        logging.info("handle_index_get")
        return web.FileResponse("./index.html")


    async def websocket_handler(request):

        global prompt_data

        ws = web.WebSocketResponse()

        await ws.prepare(request)

        logging.info("Websocket connected.")

        request.app['websockets'].add(ws)

        try:
            async for msg in ws:
                logging.info(f"Received message from websocket.")
                if "prompt" in msg.data:
                    header, prompt = msg.data.split(":")
                    logging.info("Received prompt: " + prompt)
                    try:
                        tree = Tree.from_prompt(prompt)
                        clip_encodings = predictor.encode_clip_text(tree)
                        owl_encodings = predictor.encode_owl_text(tree)
                        prompt_data = {
                            "tree": tree,
                            "clip_encodings": clip_encodings,
                            "owl_encodings": owl_encodings
                        }
                        logging.info("Set prompt: " + prompt)
                    except Exception as e:
                        print(e)
        finally:
            request.app['websockets'].discard(ws)

        return ws


    async def on_shutdown(app: web.Application):
        for ws in set(app['websockets']):
            await ws.close(code=WSCloseCode.GOING_AWAY,
                        message='Server shutdown')


    async def detection_loop(app: web.Application):

        loop = asyncio.get_running_loop()

        logging.info("Opening camera.")

        camera = cv2.VideoCapture(CAMERA_DEVICE)

        logging.info("Loading predictor.")

        def _read_and_encode_image():

            re, image = camera.read()

            if not re:
                return re, None

            image_pil = cv2_to_pil(image)

            if prompt_data is not None:
                prompt_data_local = prompt_data
                t0 = time.perf_counter_ns()
                detections = predictor.predict(
                    image_pil,
                    tree=prompt_data_local['tree'],
                    clip_text_encodings=prompt_data_local['clip_encodings'],
                    owl_text_encodings=prompt_data_local['owl_encodings']
                )
                t1 = time.perf_counter_ns()
                dt = (t1 - t0) / 1e9
                tree = prompt_data_local['tree']
                image = draw_tree_output(image, detections, prompt_data_local['tree'])

            image_jpeg = bytes(
                cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, IMAGE_QUALITY])[1]
            )

            return re, image_jpeg

        while True:

            re, image = await loop.run_in_executor(None, _read_and_encode_image)
            
            if not re:
                break
            
            for ws in app["websockets"]:
                await ws.send_bytes(image)

        camera.release()


    async def run_detection_loop(app):
        try:
            task = asyncio.create_task(detection_loop(app))
            yield
            task.cancel()
        except asyncio.CancelledError:
            pass
        finally:
            await task


    logging.basicConfig(level=logging.INFO)
    app = web.Application()
    app['websockets'] = weakref.WeakSet()
    app.router.add_get("/", handle_index_get)
    app.router.add_route("GET", "/ws", websocket_handler)
    app.on_shutdown.append(on_shutdown)
    app.cleanup_ctx.append(run_detection_loop)
    web.run_app(app, host=args.host, port=args.port)
examples/tree_demo/tree_demo.py
NVIDIA-AI-IOT/nanoowl
# 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 PIL.Image
from nanoowl.owl_predictor import OwlPredictor
from nanoowl.tree_predictor import (
    TreePredictor, Tree
)
from nanoowl.tree_drawing import (
    draw_tree_output
)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--image", type=str, default="../assets/owl_glove_small.jpg")
    parser.add_argument("--prompt", type=str, default="")
    parser.add_argument("--threshold", type=float, default=0.1)
    parser.add_argument("--output", type=str, default="../data/tree_predict_out.jpg")
    parser.add_argument("--model", type=str, default="google/owlvit-base-patch32")
    parser.add_argument("--image_encoder_engine", type=str, default="../data/owlvit_image_encoder_patch32.engine")
    args = parser.parse_args()

    predictor = TreePredictor(
        owl_predictor=OwlPredictor(
            args.model,
            image_encoder_engine=args.image_encoder_engine
        )
    )

    image = PIL.Image.open(args.image)
    tree = Tree.from_prompt(args.prompt)
    clip_text_encodings = predictor.encode_clip_text(tree)
    owl_text_encodings = predictor.encode_owl_text(tree)

    output = predictor.predict(
        image=image, 
        tree=tree,
        clip_text_encodings=clip_text_encodings,
        owl_text_encodings=owl_text_encodings,
        threshold=args.threshold
    )

    image = draw_tree_output(image, output, tree=tree, draw_text=True)

    image.save(args.output)
examples/tree_predict.py
NVIDIA-AI-IOT/nanoowl
# 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.

nanoowl/__init__.py
NVIDIA-AI-IOT/nanoowl
# 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
from .owl_predictor import OwlPredictor


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("output_path", type=str)
    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)
    args = parser.parse_args()
    
    predictor = OwlPredictor(
        model_name=args.model_name
    )

    predictor.build_image_encoder_engine(
        args.output_path,
        fp16_mode=args.fp16_mode,
        onnx_opset=args.onnx_opset
    )
nanoowl/build_image_encoder_engine.py
NVIDIA-AI-IOT/nanoowl
# 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 torch
import clip
import PIL.Image
from torchvision.ops import roi_align
from typing import List, Tuple, Optional
from dataclasses import dataclass
from .image_preprocessor import ImagePreprocessor


__all__ = [
    "ClipPredictor",
    "ClipEncodeTextOutput",
    "ClipEncodeImageOutput",
    "ClipDecodeOutput"
]


@dataclass
class ClipEncodeTextOutput:
    text_embeds: torch.Tensor

    def slice(self, start_index, end_index):
        return ClipEncodeTextOutput(
            text_embeds=self.text_embeds[start_index:end_index]
        )


@dataclass
class ClipEncodeImageOutput:
    image_embeds: torch.Tensor


@dataclass
class ClipDecodeOutput:
    labels: torch.Tensor
    scores: torch.Tensor


class ClipPredictor(torch.nn.Module):
    
    def __init__(self,
            model_name: str = "ViT-B/32",
            image_size: Tuple[int, int] = (224, 224),
            device: str = "cuda",
            image_preprocessor: Optional[ImagePreprocessor] = None
        ):
        super().__init__()
        self.device = device
        self.clip_model, _ = clip.load(model_name, device)
        self.image_size = image_size
        self.mesh_grid = torch.stack(
            torch.meshgrid(
                torch.linspace(0., 1., self.image_size[1]),
                torch.linspace(0., 1., self.image_size[0])
            )
        ).to(self.device).float()
        self.image_preprocessor = image_preprocessor.to(self.device).eval() if image_preprocessor else ImagePreprocessor().to(self.device).eval()
    
    def get_device(self):
        return self.device
        
    def get_image_size(self):
        return self.image_size

    def encode_text(self, text: List[str]) -> ClipEncodeTextOutput:
        text_tokens = clip.tokenize(text).to(self.device)
        text_embeds = self.clip_model.encode_text(text_tokens)
        return ClipEncodeTextOutput(text_embeds=text_embeds)

    def encode_image(self, image: torch.Tensor) -> ClipEncodeImageOutput:
        image_embeds = self.clip_model.encode_image(image)
        return ClipEncodeImageOutput(image_embeds=image_embeds)

    def extract_rois(self, image: torch.Tensor, rois: torch.Tensor, pad_square: bool = True, padding_scale: float=1.0):
        if len(rois) == 0:
            return torch.empty(
                (0, image.shape[1], self.image_size[0], self.image_size[1]),
                dtype=image.dtype,
                device=image.device
            )

        if pad_square:
            # pad square
            w = padding_scale * (rois[..., 2] - rois[..., 0]) / 2
            h = padding_scale * (rois[..., 3] - rois[..., 1]) / 2
            cx = (rois[..., 0] + rois[..., 2]) / 2
            cy = (rois[..., 1] + rois[..., 3]) / 2
            s = torch.max(w, h)
            rois = torch.stack([cx-s, cy-s, cx+s, cy+s], dim=-1)

            # compute mask
            pad_x = (s - w) / (2 * s)
            pad_y = (s - h) / (2 * s)
            mask_x = (self.mesh_grid[1][None, ...] > pad_x[..., None, None]) & (self.mesh_grid[1][None, ...] < (1. - pad_x[..., None, None]))
            mask_y = (self.mesh_grid[0][None, ...] > pad_y[..., None, None]) & (self.mesh_grid[0][None, ...] < (1. - pad_y[..., None, None]))
            mask = (mask_x & mask_y)

        roi_images = roi_align(image, [rois], output_size=self.get_image_size())

        if pad_square:
            roi_images = roi_images * mask[:, None, :, :]

        return roi_images, rois

    def encode_rois(self, image: torch.Tensor, rois: torch.Tensor, pad_square: bool = True, padding_scale: float = 1.0):
        roi_images, rois = self.extract_rois(image, rois, pad_square, padding_scale)
        return self.encode_image(roi_images)

    def decode(self, 
            image_output: ClipEncodeImageOutput, 
            text_output: ClipEncodeTextOutput
        ) -> ClipDecodeOutput:

        image_embeds = image_output.image_embeds
        text_embeds = text_output.text_embeds

        image_embeds = image_embeds / image_embeds.norm(dim=1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(dim=1, keepdim=True)
        logit_scale = self.clip_model.logit_scale.exp()
        logits_per_image = logit_scale * image_embeds @ text_embeds.t()
        probs = torch.softmax(logits_per_image, dim=-1)
        prob_max = probs.max(dim=-1)
        
        return ClipDecodeOutput(
            labels=prob_max.indices,
            scores=prob_max.values
        )

    def predict(self, 
            image: PIL.Image, 
            text: List[str], 
            text_encodings: Optional[ClipEncodeTextOutput],
            pad_square: bool = True,
            threshold: float = 0.1
        ) -> ClipDecodeOutput:

        image_tensor = self.image_preprocessor.preprocess_pil_image(image)

        if text_encodings is None:
            text_encodings = self.encode_text(text)

        rois = torch.tensor([[0, 0, image.height, image.width]], dtype=image_tensor.dtype, device=image_tensor.device)

        image_encodings = self.encode_rois(image_tensor, rois, pad_square=pad_square)

        return self.decode(image_encodings, text_encodings)
nanoowl/clip_predictor.py
NVIDIA-AI-IOT/nanoowl
# 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 torch
import PIL.Image
import numpy as np
from typing import Tuple


__all__ = [
    "ImagePreprocessor",
    "DEFAULT_IMAGE_PREPROCESSOR_MEAN",
    "DEFAULT_IMAGE_PREPROCESSOR_STD"
]


DEFAULT_IMAGE_PREPROCESSOR_MEAN = [
    0.48145466 * 255., 
    0.4578275 * 255., 
    0.40821073 * 255.
]


DEFAULT_IMAGE_PREPROCESSOR_STD = [
    0.26862954 * 255., 
    0.26130258 * 255., 
    0.27577711 * 255.
]


class ImagePreprocessor(torch.nn.Module):
    def __init__(self,
            mean: Tuple[float, float, float] = DEFAULT_IMAGE_PREPROCESSOR_MEAN,
            std: Tuple[float, float, float] = DEFAULT_IMAGE_PREPROCESSOR_STD
        ):
        super().__init__()
        
        self.register_buffer(
            "mean",
            torch.tensor(mean)[None, :, None, None]
        )
        self.register_buffer(
            "std",
            torch.tensor(std)[None, :, None, None]
        )

    def forward(self, image: torch.Tensor, inplace: bool = False):

        if inplace:
            image = image.sub_(self.mean).div_(self.std)
        else:
            image = (image - self.mean) / self.std

        return image
    
    @torch.no_grad()
    def preprocess_pil_image(self, image: PIL.Image.Image):
        image = torch.from_numpy(np.asarray(image))
        image = image.permute(2, 0, 1)[None, ...]
        image = image.to(self.mean.device)
        image = image.type(self.mean.dtype)
        return self.forward(image, inplace=True)
nanoowl/image_preprocessor.py
NVIDIA-AI-IOT/nanoowl
# 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 PIL.Image
import PIL.ImageDraw
import cv2
from .owl_predictor import OwlDecodeOutput
import matplotlib.pyplot as plt
import numpy as np
from typing import List


def get_colors(count: int):
    cmap = plt.cm.get_cmap("rainbow", count)
    colors = []
    for i in range(count):
        color = cmap(i)
        color = [int(255 * value) for value in color]
        colors.append(tuple(color))
    return colors


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)
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.75
    colors = get_colors(len(text))
    num_detections = len(output.labels)

    for i in range(num_detections):
        box = output.boxes[i]
        label_index = int(output.labels[i])
        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
        )
        if draw_text:
            offset_y = 12
            offset_x = 0
            label_text = text[label_index]
            cv2.putText(
                image,
                label_text,
                (box[0] + offset_x, box[1] + offset_y),
                font,
                font_scale,
                colors[label_index],
                2,# thickness
                cv2.LINE_AA
            )
    if is_pil:
        image = PIL.Image.fromarray(image)
    return image
nanoowl/owl_drawing.py
NVIDIA-AI-IOT/nanoowl
# 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 torch
import numpy as np
import PIL.Image
import subprocess
import tempfile
import os
from torchvision.ops import roi_align
from transformers.models.owlvit.modeling_owlvit import OwlViTForObjectDetection
from transformers.models.owlvit.processing_owlvit import OwlViTProcessor
from dataclasses import dataclass
from typing import List, Optional, Union, Tuple
from .image_preprocessor import ImagePreprocessor

__all__ = [
    "OwlPredictor",
    "OwlEncodeTextOutput",
    "OwlEncodeImageOutput",
    "OwlDecodeOutput"
]


def _owl_center_to_corners_format_torch(bboxes_center):
    center_x, center_y, width, height = bboxes_center.unbind(-1)
    bbox_corners = torch.stack(
        [
            (center_x - 0.5 * width), 
            (center_y - 0.5 * height), 
            (center_x + 0.5 * width), 
            (center_y + 0.5 * height)
        ],
        dim=-1,
    )
    return bbox_corners


def _owl_get_image_size(hf_name: str):

    image_sizes = {
        "google/owlvit-base-patch32": 768,
        "google/owlvit-base-patch16": 768,
        "google/owlvit-large-patch14": 840,
    }

    return image_sizes[hf_name]


def _owl_get_patch_size(hf_name: str):

    patch_sizes = {
        "google/owlvit-base-patch32": 32,
        "google/owlvit-base-patch16": 16,
        "google/owlvit-large-patch14": 14,
    }

    return patch_sizes[hf_name]


# This function is modified from https://github.com/huggingface/transformers/blob/e8fdd7875def7be59e2c9b823705fbf003163ea0/src/transformers/models/owlvit/modeling_owlvit.py#L1333
# Copyright 2022 Google AI and The HuggingFace Team. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
def _owl_normalize_grid_corner_coordinates(num_patches_per_side):
    box_coordinates = np.stack(
        np.meshgrid(np.arange(1, num_patches_per_side + 1), np.arange(1, num_patches_per_side + 1)), axis=-1
    ).astype(np.float32)
    box_coordinates /= np.array([num_patches_per_side, num_patches_per_side], np.float32)

    box_coordinates = box_coordinates.reshape(
        box_coordinates.shape[0] * box_coordinates.shape[1], box_coordinates.shape[2]
    )
    box_coordinates = torch.from_numpy(box_coordinates)

    return box_coordinates


# This function is modified from https://github.com/huggingface/transformers/blob/e8fdd7875def7be59e2c9b823705fbf003163ea0/src/transformers/models/owlvit/modeling_owlvit.py#L1354
# Copyright 2022 Google AI and The HuggingFace Team. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
def _owl_compute_box_bias(num_patches_per_side):
    box_coordinates = _owl_normalize_grid_corner_coordinates(num_patches_per_side)
    box_coordinates = torch.clip(box_coordinates, 0.0, 1.0)

    box_coord_bias = torch.log(box_coordinates + 1e-4) - torch.log1p(-box_coordinates + 1e-4)

    box_size = torch.full_like(box_coord_bias, 1.0 / num_patches_per_side)
    box_size_bias = torch.log(box_size + 1e-4) - torch.log1p(-box_size + 1e-4)

    box_bias = torch.cat([box_coord_bias, box_size_bias], dim=-1)

    return box_bias


def _owl_box_roi_to_box_global(boxes, rois):
    x0y0 = rois[..., :2]
    x1y1 = rois[..., 2:]
    wh = (x1y1 - x0y0).repeat(1, 1, 2)
    x0y0 = x0y0.repeat(1, 1, 2)
    return (boxes * wh) + x0y0


@dataclass
class OwlEncodeTextOutput:
    text_embeds: torch.Tensor

    def slice(self, start_index, end_index):
        return OwlEncodeTextOutput(
            text_embeds=self.text_embeds[start_index:end_index]
        )


@dataclass
class OwlEncodeImageOutput:
    image_embeds: torch.Tensor
    image_class_embeds: torch.Tensor
    logit_shift: torch.Tensor
    logit_scale: torch.Tensor
    pred_boxes: torch.Tensor


@dataclass
class OwlDecodeOutput:
    labels: torch.Tensor
    scores: torch.Tensor
    boxes: torch.Tensor
    input_indices: torch.Tensor


class OwlPredictor(torch.nn.Module):
    
    def __init__(self,
            model_name: str = "google/owlvit-base-patch32",
            device: str = "cuda",
            image_encoder_engine: Optional[str] = None,
            image_encoder_engine_max_batch_size: int = 1,
            image_preprocessor: Optional[ImagePreprocessor] = None
        ):

        super().__init__()

        self.image_size = _owl_get_image_size(model_name)
        self.device = device
        self.model = OwlViTForObjectDetection.from_pretrained(model_name).to(self.device).eval()
        self.processor = OwlViTProcessor.from_pretrained(model_name)
        self.patch_size = _owl_get_patch_size(model_name)
        self.num_patches_per_side = self.image_size // self.patch_size
        self.box_bias = _owl_compute_box_bias(self.num_patches_per_side).to(self.device)
        self.num_patches = (self.num_patches_per_side)**2
        self.mesh_grid = torch.stack(
            torch.meshgrid(
                torch.linspace(0., 1., self.image_size),
                torch.linspace(0., 1., self.image_size)
            )
        ).to(self.device).float()
        self.image_encoder_engine = None
        if image_encoder_engine is not None:
            image_encoder_engine = OwlPredictor.load_image_encoder_engine(image_encoder_engine, image_encoder_engine_max_batch_size)
        self.image_encoder_engine = image_encoder_engine
        self.image_preprocessor = image_preprocessor.to(self.device).eval() if image_preprocessor else ImagePreprocessor().to(self.device).eval()

    def get_num_patches(self):
        return self.num_patches

    def get_device(self):
        return self.device
        
    def get_image_size(self):
        return (self.image_size, self.image_size)
    
    def encode_text(self, text: List[str]) -> OwlEncodeTextOutput:
        text_input = self.processor(text=text, return_tensors="pt")
        input_ids = text_input['input_ids'].to(self.device)
        attention_mask = text_input['attention_mask'].to(self.device)
        text_outputs = self.model.owlvit.text_model(input_ids, attention_mask)
        text_embeds = text_outputs[1]
        text_embeds = self.model.owlvit.text_projection(text_embeds)
        return OwlEncodeTextOutput(text_embeds=text_embeds)

    def encode_image_torch(self, image: torch.Tensor) -> OwlEncodeImageOutput:
        
        vision_outputs = self.model.owlvit.vision_model(image)
        last_hidden_state = vision_outputs[0]
        image_embeds = self.model.owlvit.vision_model.post_layernorm(last_hidden_state)
        class_token_out = image_embeds[:, :1, :]
        image_embeds = image_embeds[:, 1:, :] * class_token_out
        image_embeds = self.model.layer_norm(image_embeds)  # 768 dim

        # Box Head
        pred_boxes = self.model.box_head(image_embeds)
        pred_boxes += self.box_bias
        pred_boxes = torch.sigmoid(pred_boxes)
        pred_boxes = _owl_center_to_corners_format_torch(pred_boxes)

        # Class Head
        image_class_embeds = self.model.class_head.dense0(image_embeds)
        logit_shift = self.model.class_head.logit_shift(image_embeds)
        logit_scale = self.model.class_head.logit_scale(image_embeds)
        logit_scale = self.model.class_head.elu(logit_scale) + 1

        output = OwlEncodeImageOutput(
            image_embeds=image_embeds,
            image_class_embeds=image_class_embeds,
            logit_shift=logit_shift,
            logit_scale=logit_scale,
            pred_boxes=pred_boxes
        )

        return output
    
    def encode_image_trt(self, image: torch.Tensor) -> OwlEncodeImageOutput:
        return self.image_encoder_engine(image)

    def encode_image(self, image: torch.Tensor) -> OwlEncodeImageOutput:
        if self.image_encoder_engine is not None:
            return self.encode_image_trt(image)
        else:
            return self.encode_image_torch(image)

    def extract_rois(self, image: torch.Tensor, rois: torch.Tensor, pad_square: bool = True, padding_scale: float = 1.0):
        if len(rois) == 0:
            return torch.empty(
                (0, image.shape[1], self.image_size, self.image_size),
                dtype=image.dtype,
                device=image.device
            )
        if pad_square:
            # pad square
            w = padding_scale * (rois[..., 2] - rois[..., 0]) / 2
            h = padding_scale * (rois[..., 3] - rois[..., 1]) / 2
            cx = (rois[..., 0] + rois[..., 2]) / 2
            cy = (rois[..., 1] + rois[..., 3]) / 2
            s = torch.max(w, h)
            rois = torch.stack([cx-s, cy-s, cx+s, cy+s], dim=-1)

            # compute mask
            pad_x = (s - w) / (2 * s)
            pad_y = (s - h) / (2 * s)
            mask_x = (self.mesh_grid[1][None, ...] > pad_x[..., None, None]) & (self.mesh_grid[1][None, ...] < (1. - pad_x[..., None, None]))
            mask_y = (self.mesh_grid[0][None, ...] > pad_y[..., None, None]) & (self.mesh_grid[0][None, ...] < (1. - pad_y[..., None, None]))
            mask = (mask_x & mask_y)

        # extract rois
        roi_images = roi_align(image, [rois], output_size=self.get_image_size())

        # mask rois
        if pad_square:
            roi_images = (roi_images * mask[:, None, :, :])

        return roi_images, rois
    
    def encode_rois(self, image: torch.Tensor, rois: torch.Tensor, pad_square: bool = True, padding_scale: float=1.0):
        # with torch_timeit_sync("extract rois"):
        roi_images, rois = self.extract_rois(image, rois, pad_square, padding_scale)
        # with torch_timeit_sync("encode images"):
        output = self.encode_image(roi_images)
        pred_boxes = _owl_box_roi_to_box_global(output.pred_boxes, rois[:, None, :])
        output.pred_boxes = pred_boxes
        return output

    def decode(self, 
            image_output: OwlEncodeImageOutput, 
            text_output: OwlEncodeTextOutput,
            threshold: 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)
        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)
        labels = scores_max.indices
        scores = scores_max.values

        mask = (scores > threshold)

        input_indices = torch.arange(0, num_input_images, dtype=labels.dtype, device=labels.device)
        input_indices = input_indices[:, None].repeat(1, self.num_patches)

        return OwlDecodeOutput(
            labels=labels[mask],
            scores=scores[mask],
            boxes=image_output.pred_boxes[mask],
            input_indices=input_indices[mask]
        )

    @staticmethod
    def get_image_encoder_input_names():
        return ["image"]

    @staticmethod
    def get_image_encoder_output_names():
        names = [
            "image_embeds",
            "image_class_embeds",
            "logit_shift",
            "logit_scale",
            "pred_boxes"
        ]
        return names


    def export_image_encoder_onnx(self, 
            output_path: str,
            use_dynamic_axes: bool = True,
            batch_size: int = 1,
            onnx_opset=17
        ):
        
        class TempModule(torch.nn.Module):
            def __init__(self, parent):
                super().__init__()
                self.parent = parent
            def forward(self, image):
                output = self.parent.encode_image_torch(image)
                return (
                    output.image_embeds,
                    output.image_class_embeds,
                    output.logit_shift,
                    output.logit_scale,
                    output.pred_boxes
                )

        data = torch.randn(batch_size, 3, self.image_size, self.image_size).to(self.device)

        if use_dynamic_axes:
            dynamic_axes =  {
                "image": {0: "batch"},
                "image_embeds": {0: "batch"},
                "image_class_embeds": {0: "batch"},
                "logit_shift": {0: "batch"},
                "logit_scale": {0: "batch"},
                "pred_boxes": {0: "batch"}       
            }
        else:
            dynamic_axes = {}

        model = TempModule(self)

        torch.onnx.export(
            model, 
            data, 
            output_path, 
            input_names=self.get_image_encoder_input_names(), 
            output_names=self.get_image_encoder_output_names(),
            dynamic_axes=dynamic_axes,
            opset_version=onnx_opset
        )
    
    @staticmethod
    def load_image_encoder_engine(engine_path: str, max_batch_size: int = 1):
        import tensorrt as trt
        from torch2trt import TRTModule

        with trt.Logger() as logger, trt.Runtime(logger) as runtime:
            with open(engine_path, 'rb') as f:
                engine_bytes = f.read()
            engine = runtime.deserialize_cuda_engine(engine_bytes)

        base_module = TRTModule(
            engine,
            input_names=OwlPredictor.get_image_encoder_input_names(),
            output_names=OwlPredictor.get_image_encoder_output_names()
        )

        class Wrapper(torch.nn.Module):
            def __init__(self, base_module: TRTModule, max_batch_size: int):
                super().__init__()
                self.base_module = base_module
                self.max_batch_size = max_batch_size

            @torch.no_grad()
            def forward(self, image):

                b = image.shape[0]

                results = []

                for start_index in range(0, b, self.max_batch_size):
                    end_index = min(b, start_index + self.max_batch_size)
                    image_slice = image[start_index:end_index]
                    # with torch_timeit_sync("run_engine"):
                    output = self.base_module(image_slice)
                    results.append(
                        output
                    )

                return OwlEncodeImageOutput(
                    image_embeds=torch.cat([r[0] for r in results], dim=0),
                    image_class_embeds=torch.cat([r[1] for r in results], dim=0),
                    logit_shift=torch.cat([r[2] for r in results], dim=0),
                    logit_scale=torch.cat([r[3] for r in results], dim=0),
                    pred_boxes=torch.cat([r[4] for r in results], dim=0)
                )

        image_encoder = Wrapper(base_module, max_batch_size)

        return image_encoder

    def build_image_encoder_engine(self, 
            engine_path: str, 
            max_batch_size: int = 1, 
            fp16_mode = True, 
            onnx_path: Optional[str] = None,
            onnx_opset: int = 17
        ):

        if onnx_path is None:
            onnx_dir = tempfile.mkdtemp()
            onnx_path = os.path.join(onnx_dir, "image_encoder.onnx")
            self.export_image_encoder_onnx(onnx_path, onnx_opset=onnx_opset)

        args = ["/usr/src/tensorrt/bin/trtexec"]
    
        args.append(f"--onnx={onnx_path}")
        args.append(f"--saveEngine={engine_path}")

        if fp16_mode:
            args += ["--fp16"]

        args += [f"--shapes=image:1x3x{self.image_size}x{self.image_size}"]

        subprocess.call(args)

        return self.load_image_encoder_engine(engine_path, max_batch_size)

    def predict(self, 
            image: PIL.Image, 
            text: List[str], 
            text_encodings: Optional[OwlEncodeTextOutput],
            pad_square: bool = True,
            threshold: float = 0.1
        ) -> OwlDecodeOutput:

        image_tensor = self.image_preprocessor.preprocess_pil_image(image)

        if text_encodings is None:
            text_encodings = self.encode_text(text)

        rois = torch.tensor([[0, 0, image.height, image.width]], dtype=image_tensor.dtype, device=image_tensor.device)

        image_encodings = self.encode_rois(image_tensor, rois, pad_square=pad_square)

        return self.decode(image_encodings, text_encodings, threshold)


nanoowl/owl_predictor.py
NVIDIA-AI-IOT/nanoowl
# 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 time
import torch


__all__ == [
    "SyncTimer"
]


class SyncTimer():

    def __init__(self, name: str):
        self.name = name
        self.t0 = None

    def __enter__(self, *args, **kwargs):
        self.t0 = time.perf_counter_ns()

    def __exit__(self, *args, **kwargs):
        torch.cuda.current_stream().synchronize()
        t1 = time.perf_counter_ns()
        dt = (t1 - self.t0) / 1e9
        print(f"{self.name} FPS: {round(1./dt, 3)}")
nanoowl/sync_timer.py
NVIDIA-AI-IOT/nanoowl
# 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 json
from enum import Enum
from typing import List, Optional, Mapping
from .clip_predictor import ClipEncodeTextOutput
from .owl_predictor import OwlEncodeTextOutput


__all__ = [
    "TreeOp",
    "TreeNode",
    "Tree"
]


class TreeOp(Enum):
    DETECT = "detect"
    CLASSIFY = "classify"

    def __str__(self) -> str:
        return str(self.value)


class TreeNode:
    op: TreeOp
    input: int
    outputs: List[int]

    def __init__(self, op: TreeOp, input: int, outputs: Optional[List[int]] = None):
        self.op = op
        self.input = input
        self.outputs = [] if outputs is None else outputs

    def to_dict(self):
        return {
            "op": str(self.op),
            "input": self.input,
            "outputs": self.outputs
        }

    @staticmethod
    def from_dict(node_dict: dict):

        if "op" not in node_dict:
            raise RuntimeError("Missing 'op' field.")

        if "input" not in node_dict:
            raise RuntimeError("Missing 'input' field.")
        
        if "outputs" not in node_dict:
            raise RuntimeError("Missing 'input' field.")
        
        return TreeNode(
            op=node_dict["op"],
            input=node_dict["input"],
            outputs=node_dict["outputs"]
        )
    

class Tree:
    nodes: List[TreeNode]
    labels: List[str]

    def __init__(self, nodes, labels):
        self.nodes = nodes
        self.labels = labels
        self._label_index_to_node_map = self._build_label_index_to_node_map()
    
    def _build_label_index_to_node_map(self) -> Mapping[int, "TreeNode"]:
        label_to_node_map = {}
        for node in self.nodes:
            for label_index in node.outputs:
                if label_index in label_to_node_map:
                    raise RuntimeError("Duplicate output label.")
                label_to_node_map[label_index] = node
        return label_to_node_map
    
    def to_dict(self):
        return {
            "nodes": [node.to_dict() for node in self.nodes],
            "labels": self.labels
        }

    @staticmethod
    def from_prompt(prompt: str):

        nodes = []
        node_stack = []
        label_index_stack = [0]
        labels = ["image"]
        label_index = 0

        for ch in prompt:

            if ch == "[":
                label_index += 1
                node = TreeNode(op=TreeOp.DETECT, input=label_index_stack[-1])
                node.outputs.append(label_index)
                node_stack.append(node)
                label_index_stack.append(label_index)
                labels.append("")
                nodes.append(node)
            elif ch == "]":
                if len(node_stack) == 0:
                    raise RuntimeError("Unexpected ']'.")
                node = node_stack.pop()
                if node.op != TreeOp.DETECT:
                    raise RuntimeError("Unexpected ']'.")
                label_index_stack.pop()
            elif ch == "(":
                label_index = label_index + 1
                node = TreeNode(op=TreeOp.CLASSIFY, input=label_index_stack[-1])
                node.outputs.append(label_index)
                node_stack.append(node)
                label_index_stack.append(label_index)
                labels.append("")
                nodes.append(node)
            elif ch == ")":
                if len(node_stack) == 0:
                    raise RuntimeError("Unexpected ')'.")
                node = node_stack.pop()
                if node.op != TreeOp.CLASSIFY:
                    raise RuntimeError("Unexpected ')'.")
                label_index_stack.pop()
            elif ch == ",":
                label_index_stack.pop()
                label_index = label_index + 1
                label_index_stack.append(label_index)
                node_stack[-1].outputs.append(label_index)
                labels.append("")
            else:
                labels[label_index_stack[-1]] += ch

        if len(node_stack) > 0:
            if node_stack[-1].op == TreeOp.DETECT:
                raise RuntimeError("Missing ']'.")
            if node_stack[-1].op == TreeOp.CLASSIFY:
                raise RuntimeError("Missing ')'.")
            
        labels = [label.strip() for label in labels]

        graph = Tree(nodes=nodes, labels=labels)

        return graph
    
    def to_json(self, indent: Optional[int] = None) -> str:
        return json.dumps(self.to_dict(), indent=indent)
    
    @staticmethod
    def from_dict(tree_dict: dict) -> "Tree":

        if "nodes" not in tree_dict:
            raise RuntimeError("Missing 'nodes' field.")
        
        if "labels" not in tree_dict:
            raise RuntimeError("Missing 'labels' field.")
        
        nodes = [TreeNode.from_dict(node_dict) for node_dict in tree_dict["nodes"]]
        labels = tree_dict["labels"]

        return Tree(nodes=nodes, labels=labels)

    @staticmethod
    def from_json(tree_json: str) -> "Tree":
        tree_dict = json.loads(tree_json)
        return Tree.from_dict(tree_dict)
    
    def get_op_for_label_index(self, label_index: int):
        if label_index not in self._label_index_to_node_map:
            return None
        return self._label_index_to_node_map[label_index].op
    
    def get_label_indices_with_op(self, op: TreeOp):
        return [
            index for index in range(len(self.labels))
            if self.get_op_for_label_index(index) == op
        ]
    
    def get_classify_label_indices(self):
        return self.get_label_indices_with_op(TreeOp.CLASSIFY)
    
    def get_detect_label_indices(self):
        return self.get_label_indices_with_op(TreeOp.DETECT)

    def find_nodes_with_input(self, input_index: int):
        return [n for n in self.nodes if n.input == input_index]

    def find_detect_nodes_with_input(self, input_index: int):
        return [n for n in self.find_nodes_with_input(input_index) if n.op == TreeOp.DETECT]

    def find_classify_nodes_with_input(self, input_index: int):
        return [n for n in self.find_nodes_with_input(input_index) if n.op == TreeOp.CLASSIFY]

    def get_label_depth(self, index):
        depth = 0
        while index in self._label_index_to_node_map:
            depth += 1
            node = self._label_index_to_node_map[index]
            index = node.input
        return depth

    def get_label_depth_map(self):
        depths = {}
        for i in range(len(self.labels)):
            depths[i] = self.get_label_depth(i)
        return depths

    def get_label_map(self):
        label_map = {}
        for i in range(len(self.labels)):
            label_map[i] = self.labels[i]
        return label_map
nanoowl/tree.py
NVIDIA-AI-IOT/nanoowl
# 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 PIL.Image
import PIL.ImageDraw
import cv2
from .tree import Tree
from .tree_predictor import TreeOutput
import matplotlib.pyplot as plt
import numpy as np
from typing import List


def get_colors(count: int):
    cmap = plt.cm.get_cmap("rainbow", count)
    colors = []
    for i in range(count):
        color = cmap(i)
        color = [int(255 * value) for value in color]
        colors.append(tuple(color))
    return colors


def draw_tree_output(image, output: TreeOutput, tree: Tree, draw_text=True, num_colors=8):
    detections = output.detections
    is_pil = not isinstance(image, np.ndarray)
    if is_pil:
        image = np.asarray(image)
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.75
    colors = get_colors(num_colors)
    label_map = tree.get_label_map()
    label_depths = tree.get_label_depth_map()
    for detection in detections:
        box = [int(x) for x in detection.box]
        pt0 = (box[0], box[1])
        pt1 = (box[2], box[3])
        box_depth = min(label_depths[i] for i in detection.labels)
        cv2.rectangle(
            image,
            pt0,
            pt1,
            colors[box_depth % num_colors],
            3
        )
        if draw_text:
            offset_y = 30
            offset_x = 8
            for label in detection.labels:
                label_text = label_map[label]
                cv2.putText(
                    image,
                    label_text,
                    (box[0] + offset_x, box[1] + offset_y),
                    font,
                    font_scale,
                    colors[label % num_colors],
                    2,# thickness
                    cv2.LINE_AA
                )
                offset_y += 18
    if is_pil:
        image = PIL.Image.fromarray(image)
    return image
nanoowl/tree_drawing.py
NVIDIA-AI-IOT/nanoowl
# 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 .tree import Tree, TreeOp
from .owl_predictor import OwlPredictor, OwlEncodeTextOutput, OwlEncodeImageOutput
from .clip_predictor import ClipPredictor, ClipEncodeTextOutput, ClipEncodeImageOutput
from .image_preprocessor import ImagePreprocessor

import torch
import PIL.Image
from typing import Optional, Tuple, List, Mapping, Dict
from dataclasses import dataclass


@dataclass
class TreeDetection:
    id: int
    parent_id: int
    box: Tuple[float, float, float, float]
    labels: List[int]
    scores: List[int]


@dataclass
class TreeOutput:
    detections: List[TreeDetection]
    

class TreePredictor(torch.nn.Module):

    def __init__(self,
            owl_predictor: Optional[OwlPredictor] = None,
            clip_predictor: Optional[ClipPredictor] = None,
            image_preprocessor: Optional[ImagePreprocessor] = None,
            device: str = "cuda"
        ):
        super().__init__()
        self.owl_predictor = OwlPredictor() if owl_predictor is None else owl_predictor
        self.clip_predictor = ClipPredictor() if clip_predictor is None else clip_predictor
        self.image_preprocessor = ImagePreprocessor().to(device).eval() if image_preprocessor is None else image_preprocessor

    def encode_clip_text(self, tree: Tree) -> Dict[int, ClipEncodeTextOutput]:
        label_indices = tree.get_classify_label_indices()
        if len(label_indices) == 0:
            return {}
        labels = [tree.labels[index] for index in label_indices]
        text_encodings = self.clip_predictor.encode_text(labels)
        label_encodings = {}
        for i in range(len(labels)):
            label_encodings[label_indices[i]] = text_encodings.slice(i, i+1)
        return label_encodings
    
    def encode_owl_text(self, tree: Tree) -> Dict[int, OwlEncodeTextOutput]:
        label_indices = tree.get_detect_label_indices()
        if len(label_indices) == 0:
            return {}
        labels = [tree.labels[index] for index in label_indices]
        text_encodings = self.owl_predictor.encode_text(labels)
        label_encodings = {}
        for i in range(len(labels)):
            label_encodings[label_indices[i]] = text_encodings.slice(i, i+1)
        return label_encodings
    
    @torch.no_grad()
    def predict(self, 
            image: PIL.Image.Image, 
            tree: Tree, 
            threshold: float = 0.1,
            clip_text_encodings: Optional[Dict[int, ClipEncodeTextOutput]] = None,
            owl_text_encodings: Optional[Dict[int, OwlEncodeTextOutput]] = None
        ):

        if clip_text_encodings is None:
            clip_text_encodings = self.encode_clip_text(tree)
        
        if owl_text_encodings is None:
            owl_text_encodings = self.encode_owl_text(tree)
        
        image_tensor = self.image_preprocessor.preprocess_pil_image(image)
        boxes = {
            0: torch.tensor([[0, 0, image.width, image.height]], dtype=image_tensor.dtype, device=image_tensor.device)
        }
        scores = {
            0: torch.tensor([1.], dtype=torch.float, device=image_tensor.device)
        }
        instance_ids = {
            0: torch.tensor([0], dtype=torch.int64, device=image_tensor.device)
        }
        parent_instance_ids = {
            0: torch.tensor([-1], dtype=torch.int64, device=image_tensor.device)
        }

        owl_image_encodings: Dict[int, OwlEncodeImageOutput] = {}
        clip_image_encodings: Dict[int, ClipEncodeImageOutput] = {}

        global_instance_id = 1

        queue = [0]

        while queue:
            label_index = queue.pop(0)

            detect_nodes = tree.find_detect_nodes_with_input(label_index)
            classify_nodes = tree.find_classify_nodes_with_input(label_index)

            # Run OWL image encode if required
            if len(detect_nodes) > 0 and label_index not in owl_image_encodings:
                owl_image_encodings[label_index] = self.owl_predictor.encode_rois(image_tensor, boxes[label_index])
                

            # Run CLIP image encode if required
            if len(classify_nodes) > 0 and label_index not in clip_image_encodings:
                clip_image_encodings[label_index] = self.clip_predictor.encode_rois(image_tensor, boxes[label_index])

            # Decode detect nodes
            for node in detect_nodes:

                if node.input not in owl_image_encodings:
                    raise RuntimeError("Missing owl image encodings for node.")

                # gather encodings
                owl_text_encodings_for_node = OwlEncodeTextOutput(
                    text_embeds=torch.cat([
                        owl_text_encodings[i].text_embeds for i in node.outputs
                    ], dim=0)
                )
                
                owl_node_output = self.owl_predictor.decode(
                    owl_image_encodings[node.input], 
                    owl_text_encodings_for_node, 
                    threshold=threshold
                )

                num_detections = len(owl_node_output.labels)
                instance_ids_for_node = torch.arange(global_instance_id, global_instance_id + num_detections, dtype=torch.int64, device=owl_node_output.labels.device)
                parent_instance_ids_for_node = instance_ids[node.input][owl_node_output.input_indices]
                global_instance_id += num_detections

                for i in range(len(node.outputs)):
                    mask = owl_node_output.labels == i
                    out_idx = node.outputs[i]
                    boxes[out_idx] = owl_node_output.boxes[mask]
                    scores[out_idx] = owl_node_output.scores[mask]
                    instance_ids[out_idx] = instance_ids_for_node[mask]
                    parent_instance_ids[out_idx] = parent_instance_ids_for_node[mask]

            for node in classify_nodes:

                if node.input not in clip_image_encodings:
                    raise RuntimeError("Missing clip image encodings for node.")

                clip_text_encodings_for_node = ClipEncodeTextOutput(
                    text_embeds=torch.cat([
                        clip_text_encodings[i].text_embeds for i in node.outputs
                    ], dim=0)
                )

                clip_node_output = self.clip_predictor.decode(
                    clip_image_encodings[node.input], 
                    clip_text_encodings_for_node
                )

                parent_instance_ids_for_node = instance_ids[node.input]

                for i in range(len(node.outputs)):
                    mask = clip_node_output.labels == i
                    output_buffer = node.outputs[i]
                    scores[output_buffer] = clip_node_output.scores[mask].float()
                    boxes[output_buffer] = boxes[label_index][mask].float()
                    instance_ids[output_buffer] = instance_ids[node.input][mask]
                    parent_instance_ids[output_buffer] = parent_instance_ids[node.input][mask]

            for node in detect_nodes:
                for buf in node.outputs:
                    if buf in scores and len(scores[buf]) > 0:
                        queue.append(buf)

            for node in classify_nodes:
                for buf in node.outputs:
                    if buf in scores and len(scores[buf]) > 0:
                        queue.append(buf)

        # Fill outputs
        detections: Dict[int, TreeDetection] = {}
        for i in boxes.keys():
            for box, score, instance_id, parent_instance_id in zip(boxes[i], scores[i], instance_ids[i], parent_instance_ids[i]):
                instance_id = int(instance_id)
                score = float(score)
                box = box.tolist()
                parent_instance_id = int(parent_instance_id)
                if instance_id in detections:
                    detections[instance_id].labels.append(i)
                    detections[instance_id].scores.append(score)
                else:
                    detections[instance_id] = TreeDetection(
                        id=instance_id,
                        parent_id=parent_instance_id,
                        box=box,
                        labels=[i],
                        scores=[score]
                    )

        return TreeOutput(detections=detections.values())
nanoowl/tree_predictor.py
NVIDIA-AI-IOT/nanoowl
from setuptools import find_packages, setup


setup(
    name="nanoowl",
    version="0.0.0",
    packages=find_packages()
)
setup.py
NVIDIA-AI-IOT/nanoowl
# 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.

test/__init__.py
NVIDIA-AI-IOT/nanoowl
# 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 pytest
import torch
import PIL.Image
from nanoowl.clip_predictor import ClipPredictor
from nanoowl.image_preprocessor import ImagePreprocessor


def test_get_image_size():
    clip_predictor = ClipPredictor()
    assert clip_predictor.get_image_size() == (224, 224)


def test_clip_encode_text():

    clip_predictor = ClipPredictor()

    text_encode_output = clip_predictor.encode_text(["a frog", "a dog"])

    assert text_encode_output.text_embeds.shape == (2, 512)


def test_clip_encode_image():

    clip_predictor = ClipPredictor()

    image = PIL.Image.open("assets/owl_glove_small.jpg")

    image = image.resize((224, 224))

    image_preprocessor = ImagePreprocessor().to(clip_predictor.device).eval()

    image_tensor = image_preprocessor.preprocess_pil_image(image)

    image_encode_output = clip_predictor.encode_image(image_tensor)

    assert image_encode_output.image_embeds.shape == (1, 512)


def test_clip_classify():

    clip_predictor = ClipPredictor()

    image = PIL.Image.open("assets/frog.jpg")

    image = image.resize((224, 224))
    image_preprocessor = ImagePreprocessor().to(clip_predictor.device).eval()

    image_tensor = image_preprocessor.preprocess_pil_image(image)

    text_output = clip_predictor.encode_text(["a frog", "an owl"])
    image_output = clip_predictor.encode_image(image_tensor)

    classify_output = clip_predictor.decode(image_output, text_output)

    assert classify_output.labels[0] == 0
test/test_clip_predictor.py
NVIDIA-AI-IOT/nanoowl
# 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 pytest
import torch
import PIL.Image
from nanoowl.image_preprocessor import ImagePreprocessor


def test_image_preprocessor_preprocess_pil_image():

    image_preproc = ImagePreprocessor().to("cuda").eval()

    image = PIL.Image.open("assets/owl_glove_small.jpg")

    image_tensor = image_preproc.preprocess_pil_image(image)

    assert image_tensor.shape == (1, 3, 499, 756)
    assert torch.allclose(image_tensor.mean(), torch.zeros_like(image_tensor), atol=1, rtol=1)

test/test_image_preprocessor.py
NVIDIA-AI-IOT/nanoowl
# 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 pytest
import torch
import PIL.Image
from nanoowl.owl_predictor import OwlPredictor
from nanoowl.image_preprocessor import ImagePreprocessor


def test_owl_predictor_get_image_size():
    owl_predictor = OwlPredictor()
    assert owl_predictor.get_image_size() == (768, 768)


def test_owl_predictor_encode_text():

    owl_predictor = OwlPredictor()

    text_encode_output = owl_predictor.encode_text(["a frog", "a dog"])

    assert text_encode_output.text_embeds.shape == (2, 512)


def test_owl_predictor_encode_image():

    owl_predictor = OwlPredictor()

    image = PIL.Image.open("assets/owl_glove_small.jpg")

    image = image.resize(owl_predictor.get_image_size())

    image_preprocessor = ImagePreprocessor().to(owl_predictor.device).eval()

    image_tensor = image_preprocessor.preprocess_pil_image(image)

    image_encode_output = owl_predictor.encode_image(image_tensor)

    assert image_encode_output.image_class_embeds.shape == (1, owl_predictor.get_num_patches(), 512)


def test_owl_predictor_decode():

    owl_predictor = OwlPredictor()

    image = PIL.Image.open("assets/owl_glove_small.jpg")

    image = image.resize(owl_predictor.get_image_size())
    image_preprocessor = ImagePreprocessor().to(owl_predictor.device).eval()

    image_tensor = image_preprocessor.preprocess_pil_image(image)

    text_output = owl_predictor.encode_text(["an owl"])
    image_output = owl_predictor.encode_image(image_tensor)

    classify_output = owl_predictor.decode(image_output, text_output)

    assert len(classify_output.labels == 1)
    assert classify_output.labels[0] == 0
    assert classify_output.boxes.shape == (1, 4)
    assert classify_output.input_indices == 0


def test_owl_predictor_decode_multiple_images():

    owl_predictor = OwlPredictor()
    image_preprocessor = ImagePreprocessor().to(owl_predictor.device).eval()

    image_paths = [
        "assets/owl_glove_small.jpg",
        "assets/frog.jpg"
    ]

    images = []
    for image_path in image_paths:
        image = PIL.Image.open(image_path)
        image = image.resize(owl_predictor.get_image_size())
        image = image_preprocessor.preprocess_pil_image(image)
        images.append(image)

    images = torch.cat(images, dim=0)

    text_output = owl_predictor.encode_text(["an owl", "a frog"])
    image_output = owl_predictor.encode_image(images)

    decode_output = owl_predictor.decode(image_output, text_output)

    # check number of detections
    assert len(decode_output.labels == 2)

    # check owl detection
    assert decode_output.labels[0] == 0
    assert decode_output.input_indices[0] == 0
    
    # check frog detection
    assert decode_output.labels[1] == 1
    assert decode_output.input_indices[1] == 1

test/test_owl_predictor.py
NVIDIA-AI-IOT/nanoowl
# 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 pytest
from nanoowl.tree import (
    Tree,
    TreeOp
)


def test_tree_from_prompt():

    graph = Tree.from_prompt("[a face]")

    assert len(graph.nodes) == 1
    assert len(graph.labels) == 2
    assert graph.labels[0] == "image"
    assert graph.labels[1] == "a face"
    assert graph.nodes[0].op == TreeOp.DETECT

    graph = Tree.from_prompt("[a face](a dog, a cat)")

    assert len(graph.nodes) == 2
    assert len(graph.labels) == 4
    assert graph.labels[0] == "image"
    assert graph.labels[1] == "a face"
    assert graph.labels[2] == "a dog"
    assert graph.labels[3] == "a cat"
    assert graph.nodes[0].op == TreeOp.DETECT
    assert graph.nodes[1].op == TreeOp.CLASSIFY

    with pytest.raises(RuntimeError):
        Tree.from_prompt("]a face]")

    with pytest.raises(RuntimeError):
        Tree.from_prompt("[a face")

    with pytest.raises(RuntimeError):
        Tree.from_prompt("[a face)")

    with pytest.raises(RuntimeError):
        Tree.from_prompt("[a face]]")


def test_tree_to_dict():
    
    tree = Tree.from_prompt("[a[b,c(d,e)]]")
    tree_dict = tree.to_dict()
    assert "nodes" in tree_dict
    assert "labels" in tree_dict
    assert len(tree_dict["nodes"]) == 3
    assert len(tree_dict["labels"]) == 6



def test_tree_from_prompt():

    tree = Tree.from_prompt("(office, home, outdoors, gym)")

    print(tree)
test/test_tree.py
NVIDIA-AI-IOT/nanoowl
# 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 pytest
import PIL.Image
from nanoowl.tree_predictor import TreePredictor
from nanoowl.tree import Tree


def test_encode_clip_labels():
    
    predictor = TreePredictor()
    tree = Tree.from_prompt("(sunny, rainy)")

    text_encodings = predictor.encode_clip_text(tree)

    assert len(text_encodings) == 2
    assert 1 in text_encodings
    assert 2 in text_encodings
    assert text_encodings[1].text_embeds.shape == (1, 512)


def test_encode_owl_labels():
    
    predictor = TreePredictor()
    tree = Tree.from_prompt("[a face [an eye, a nose]]")

    text_encodings = predictor.encode_owl_text(tree)

    assert len(text_encodings) == 3
    assert 1 in text_encodings
    assert 2 in text_encodings
    assert 3 in text_encodings
    assert text_encodings[1].text_embeds.shape == (1, 512)


def test_encode_clip_owl_labels_mixed():
    
    predictor = TreePredictor()
    tree = Tree.from_prompt("[a face [an eye, a nose](happy, sad)]")

    owl_text_encodings = predictor.encode_owl_text(tree)
    clip_text_encodings = predictor.encode_clip_text(tree)

    assert len(owl_text_encodings) == 3
    assert len(clip_text_encodings) == 2


def test_tree_predictor_predict():

    predictor = TreePredictor()
    tree = Tree.from_prompt("[an owl]")


    image = PIL.Image.open("assets/owl_glove.jpg")

    detections = predictor.predict(image, tree)

    
def test_tree_predictor_predict():

    predictor = TreePredictor()
    tree = Tree.from_prompt("(outdoors, indoors)")


    image = PIL.Image.open("assets/owl_glove.jpg")

    detections = predictor.predict(image, tree)

    print(detections)
    
test/test_tree_predictor.py
NVIDIA-AI-IOT/nanoowl