adding models

infer.py

@@ -0,0 +1,305 @@
+import sys
+rt_path = "/home/thebears/Source/task_runners/vision_v3/cuda_objdet_clip"
+sys.path.insert(0, rt_path)
+import torch
+cc = torch.cuda.get_device_properties(0)
+cc_str = str(cc.major)+'.'+str(cc.minor)
+import os
+import tensorrt as trt
+import numpy as np
+import time
+from loaders.video_loaders import decoder
+import torch
+from torchvision.transforms import v2
+from transforms.model_transforms import det_transforms, clip_transforms
+from loaders.model_loaders import TensorRTModel
+from common_code import file_names
+import logging
+
+log = logging.getLogger()
+
+clip_frame_skip_interval = 8
+clip_cropped_frame_skip_interval = 24
+det_frame_skip_interval = 2
+
+det_threshold = 0.5
+
+
+
+clip_engine_path = os.path.join(
+    rt_path, "models/ViT-SO400M-16-SigLIP2-512_visual_fp16_"+cc_str+".engine"
+)
+obj_engine_path = os.path.join(rt_path, "models/dfine_large_"+cc_str+".engine")
+species_list_file = os.path.join(rt_path, "models/dfine_large.species_keep")
+abspath_species_list_file = os.path.abspath(species_list_file)
+with open(species_list_file, "r") as ff:
+    species_list = ff.read().split("\n")
+
+
+det_model = TensorRTModel(obj_engine_path)
+det_model.prepare()
+det_model.allocate_outputs()
+
+clip_model = TensorRTModel(clip_engine_path)
+clip_model.prepare()
+clip_model.allocate_outputs()
+
+
+
+
+def score_video_cached(file_path, return_dict = False):
+    file_path = file_names.resolve_file_location(file_path)
+    det_path = file_names.get_det_npz_path(file_path)
+    emb_path = file_names.get_embed_npz_path(file_path)
+    vc = [det_path, emb_path]
+    do_score = True
+    if all([os.path.exists(x) for x in vc]):
+        do_score = False
+        if return_dict:
+            det_dict = dict(np.load(det_path, allow_pickle = True))
+            emb_dict = dict(np.load(emb_path, allow_pickle = True))
+            det_dict['src_path'] = det_dict['src_path'].item()
+            det_dict['transform_stats'] = det_dict['transform_stats'].item()
+            emb_dict['src_path'] = emb_dict['src_path'].item()
+            n_unq_hash = len(torch.tensor(emb_dict['embeds']).hash_tensor(dim=1).unique())
+            n_total_vec = emb_dict['embeds'].shape[0]
+            if n_total_vec / n_unq_hash > 1.5:
+                log.error(f'Recreating for {emb_path} because of cyclic values')
+                do_score = True
+
+    if not do_score and return_dict:
+        return det_dict, emb_dict
+
+    if not do_score and not return_dict:
+        return
+
+    return score_video(file_path)
+
+
+def score_video_sub_clip_cached(enc_file_path):
+    crop_embeds_path = file_names.get_cropped_embed_npz_path(enc_file_path)
+    if os.path.exists(crop_embeds_path):
+        return
+    else:
+        score_video_sub_clip(enc_file_path)
+
+def score_video_sub_clip(enc_file_path):
+    crop_embeds_path = file_names.get_cropped_embed_npz_path(enc_file_path)
+    enc_file_path = file_names.resolve_file_location(enc_file_path)
+
+    nvc_batch_size = 32
+    decoder.reconfigure_decoder(enc_file_path)
+
+    all_sc = list()
+
+    keep_reading_video = True
+
+    clip_src_tensor_list = []
+    det_src_tensor_list = []
+
+    c_frm_num = 0
+
+    st = time.time()
+    n_det_scores = 0
+    n_clip_scores = 0
+
+    transform_stats = None
+    thresh_frames = list()
+    thresh_scores = list()
+    thresh_labels = list()
+    thresh_boxes = list()
+    det_all_scored_frames = list()
+    clip_embeddings = list()
+    clip_frames = list()
+    while keep_reading_video:
+        frames = decoder.get_batch_frames(nvc_batch_size)
+        if len(frames) == 0:
+            keep_reading_video = False
+
+        for frame in frames:
+            if (c_frm_num % clip_cropped_frame_skip_interval) == 0:
+                clip_src_tensor_list.append(
+                    {"tensor": torch.from_dlpack(frame), "frame_number": c_frm_num}
+                )
+            c_frm_num += 1
+
+        while len(clip_src_tensor_list) >= clip_model.batch_size:
+            clip_tens_pass = clip_src_tensor_list[0 : clip_model.batch_size]
+            clip_src_tensor_list = clip_src_tensor_list[clip_model.batch_size :]     
+            clip_frame_numbers = [x["frame_number"] for x in clip_tens_pass]
+            clip_frames.append(clip_frame_numbers)
+            clip_stacked = (
+                torch.stack([x["tensor"] for x in clip_tens_pass], dim=0) / 255.0
+            )
+
+            movie_size = list(clip_stacked.shape[::-1][0:2])
+            crop_div = 4
+            olap_factor = 0.5
+
+            crop_width = int(movie_size[0]/crop_div)
+            crop_height = int(crop_width)
+
+            spacing_width = int(crop_width * (1-olap_factor))
+            spacing_height = int(crop_height * (1-olap_factor))
+
+            starts_width = list(range(0,movie_size[0], spacing_width))
+            starts_height = list(range(0, movie_size[1], spacing_height))
+            crop_areas = set()
+            for st_w in starts_width:
+                for st_h in starts_height:
+                    en_w = st_w + crop_width
+                    en_h = st_h + crop_height
+
+                    if en_w > movie_size[0]:
+                        st_w = movie_size[0] - crop_width
+                        en_w = st_w + crop_width
+                    if en_h > movie_size[1]:
+                        st_h = movie_size[1] - crop_height
+                        en_h = st_h + crop_height            
+
+                    crop_area = ( st_w, st_h, en_w, en_h)
+                    crop_areas.add(crop_area)
+
+            cropped_embeddings = dict()
+            for crop_area in crop_areas:
+                ( st_w, st_h, en_w, en_h) = crop_area
+                clip_stack_score = clip_stacked[:,:, st_h:en_h, st_w:en_w]
+                clip_res = clip_model.score(clip_transforms(clip_stack_score).data_ptr())
+                cropped_embeddings[crop_area] = np.copy(clip_res["projected"])
+                n_clip_scores += clip_res["projected"].shape[0]
+            clip_embeddings.append(cropped_embeddings)
+
+
+    clip_dict = dict()
+    clip_dict["src_path"] = enc_file_path
+
+    cropped_embeddings = dict()
+    for c_clip in clip_embeddings:
+        for k,v in c_clip.items():
+            if k not in cropped_embeddings:
+                cropped_embeddings[k] = list()
+
+            cropped_embeddings[k].append(v)
+
+    clip_dict["frame_numbers"]  =np.concatenate(clip_frames)
+    clip_dict["embeds_cropped"] =  {k:np.concatenate(v) for k,v in cropped_embeddings.items()}
+    np.savez(crop_embeds_path, **clip_dict)
+
+
+def score_video(enc_file_path):
+    enc_file_path = file_names.resolve_file_location(enc_file_path)
+
+    nvc_batch_size = 32
+    decoder.reconfigure_decoder(enc_file_path)
+
+    all_sc = list()
+
+    keep_reading_video = True
+
+    clip_src_tensor_list = []
+    det_src_tensor_list = []
+
+    c_frm_num = 0
+
+    st = time.time()
+    n_det_scores = 0
+    n_clip_scores = 0
+
+    transform_stats = None
+    thresh_frames = list()
+    thresh_scores = list()
+    thresh_labels = list()
+    thresh_boxes = list()
+    det_all_scored_frames = list()
+    clip_embeddings = list()
+    clip_frames = list()
+    while keep_reading_video:
+
+        frames = decoder.get_batch_frames(nvc_batch_size)
+        if len(frames) == 0:
+            keep_reading_video = False
+
+        for frame in frames:
+
+            if (c_frm_num % clip_frame_skip_interval) == 0:
+                clip_src_tensor_list.append(
+                    {"tensor": torch.from_dlpack(frame), "frame_number": c_frm_num}
+                )
+
+            if (c_frm_num % det_frame_skip_interval) == 0:
+                det_src_tensor_list.append(
+                    {"tensor": torch.from_dlpack(frame), "frame_number": c_frm_num}
+                )
+
+            c_frm_num += 1
+            
+        while len(det_src_tensor_list) >= det_model.batch_size:
+
+            det_tens_pass = det_src_tensor_list[0 : det_model.batch_size]
+            det_src_tensor_list = det_src_tensor_list[det_model.batch_size :]
+            det_frame_numbers = [x["frame_number"] for x in det_tens_pass]
+            det_stacked = (
+                torch.stack([x["tensor"] for x in det_tens_pass], dim=0) / 255.0
+            )
+            det_res = det_model.score(det_transforms(det_stacked).data_ptr())
+            n_det_scores += det_res["scores"].shape[0]
+            if transform_stats is None:
+                transform_stats = det_transforms.transforms[0].transform_stats
+
+            det_scores = det_res["scores"]
+            det_labels = det_res["labels"]
+            det_boxes = det_res["boxes"]
+            det_all_scored_frames.append(det_frame_numbers)
+            (n_batch, n_queries, n_regressed) = det_boxes.shape
+
+            frames = np.repeat(
+                np.asarray(det_frame_numbers)[:, None], n_queries, axis=1
+            )
+            flattened_boxes = det_boxes.reshape((n_batch * n_queries, n_regressed))
+
+            mask_to_keep = det_scores.ravel() > det_threshold
+            
+            thresh_frames.append(frames.ravel()[mask_to_keep])
+            thresh_scores.append(det_scores.ravel()[mask_to_keep])
+            thresh_labels.append(det_labels.ravel()[mask_to_keep])
+            thresh_boxes.append(flattened_boxes[mask_to_keep, :])
+
+        while len(clip_src_tensor_list) >= clip_model.batch_size:
+ 
+            clip_tens_pass = clip_src_tensor_list[0 : clip_model.batch_size]
+            clip_src_tensor_list = clip_src_tensor_list[clip_model.batch_size :]
+            clip_frame_numbers = [x["frame_number"] for x in clip_tens_pass]
+            clip_frames.append(clip_frame_numbers)
+            clip_stacked = (
+                torch.stack([x["tensor"] for x in clip_tens_pass], dim=0) / 255.0
+            )
+            clip_res = clip_model.score(clip_transforms(clip_stacked).data_ptr())
+            clip_embeddings.append(np.copy(clip_res["projected"]))
+            n_clip_scores += clip_res["projected"].shape[0]
+
+
+
+    det_dict = dict()
+    det_dict["src_path"] = enc_file_path
+    det_dict["transform_stats"] = transform_stats
+    det_dict["scored_frames"] = np.concatenate(det_all_scored_frames)
+    det_dict["final_frames"] = np.concatenate(thresh_frames)
+    det_dict["final_labels"] = np.concatenate(thresh_labels)
+
+    det_dict["species_label_map"] = {
+        int(x): species_list[x] for x in np.unique(det_dict["final_labels"])
+    }
+    det_dict["final_scores"] = np.concatenate(thresh_scores)
+    det_dict["final_boxes"] = np.concatenate(thresh_boxes)
+
+    clip_dict = dict()
+    clip_dict["src_path"] = enc_file_path
+    clip_dict["embeds"] = np.concatenate(clip_embeddings).astype(np.float16)
+    clip_dict["frame_numbers"] = np.concatenate(clip_frames)
+
+
+    np.savez(file_names.get_embed_npz_path(enc_file_path), **clip_dict)
+    np.savez(file_names.get_det_npz_path(enc_file_path), **det_dict)
+    return det_dict, clip_dict
+
+

infer_sub_clip.py

@@ -0,0 +1,109 @@
+import sys, os
+sys.path.append('/home/thebears/Source/task_runners/vision_v3/')
+os.environ['CUDA_VISIBLE_DEVICES'] = '1'
+from cuda_objdet_clip.infer import *
+enc_file_path = '/home/thebears/Source/ml/clip_test/kya_tail/Leopards1_00_20260226160534.mp4'
+# %%
+
+crop_embeds_path = file_names.get_cropped_embed_npz_path(enc_file_path)
+enc_file_path = file_names.resolve_file_location(enc_file_path)
+
+nvc_batch_size = 32
+decoder.reconfigure_decoder(enc_file_path)
+
+all_sc = list()
+
+keep_reading_video = True
+
+clip_src_tensor_list = []
+det_src_tensor_list = []
+
+c_frm_num = 0
+
+st = time.time()
+n_det_scores = 0
+n_clip_scores = 0
+
+transform_stats = None
+thresh_frames = list()
+thresh_scores = list()
+thresh_labels = list()
+thresh_boxes = list()
+det_all_scored_frames = list()
+clip_embeddings = list()
+clip_frames = list()
+while keep_reading_video:
+    frames = decoder.get_batch_frames(nvc_batch_size)
+    if len(frames) == 0:
+        keep_reading_video = False
+
+    for frame in frames:
+        if (c_frm_num % clip_frame_skip_interval) == 0:
+            clip_src_tensor_list.append(
+                {"tensor": torch.from_dlpack(frame), "frame_number": c_frm_num}
+            )
+        c_frm_num += 1
+
+    while len(clip_src_tensor_list) >= clip_model.batch_size:
+        print('bump')
+        clip_tens_pass = clip_src_tensor_list[0 : clip_model.batch_size]
+        clip_src_tensor_list = clip_src_tensor_list[clip_model.batch_size :]     
+        clip_frame_numbers = [x["frame_number"] for x in clip_tens_pass]
+        clip_frames.append(clip_frame_numbers)
+        clip_stacked = (
+            torch.stack([x["tensor"] for x in clip_tens_pass], dim=0) / 255.0
+        )
+
+        movie_size = list(clip_stacked.shape[::-1][0:2])
+        crop_div = 4
+        olap_factor = 0.5
+        
+        crop_width = int(movie_size[0]/crop_div)
+        crop_height = int(crop_width)
+        
+        spacing_width = int(crop_width * (1-olap_factor))
+        spacing_height = int(crop_height * (1-olap_factor))
+
+        starts_width = list(range(0,movie_size[0], spacing_width))
+        starts_height = list(range(0, movie_size[1], spacing_height))
+        crop_areas = set()
+        for st_w in starts_width:
+            for st_h in starts_height:
+                en_w = st_w + crop_width
+                en_h = st_h + crop_height
+
+                if en_w > movie_size[0]:
+                    st_w = movie_size[0] - crop_width
+                    en_w = st_w + crop_width
+                if en_h > movie_size[1]:
+                    st_h = movie_size[1] - crop_height
+                    en_h = st_h + crop_height            
+
+                crop_area = ( st_w, st_h, en_w, en_h)
+                crop_areas.add(crop_area)
+
+        cropped_embeddings = dict()
+        for crop_area in crop_areas:
+            ( st_w, st_h, en_w, en_h) = crop_area
+            clip_stack_score = clip_stacked[:,:, st_h:en_h, st_w:en_w]
+            clip_res = clip_model.score(clip_transforms(clip_stack_score).data_ptr())
+            cropped_embeddings[crop_area] = np.copy(clip_res["projected"])
+            n_clip_scores += clip_res["projected"].shape[0]
+        clip_embeddings.append(cropped_embeddings)
+
+
+clip_dict = dict()
+clip_dict["src_path"] = enc_file_path
+
+cropped_embeddings = dict()
+for c_clip in clip_embeddings:
+    for k,v in c_clip.items():
+        if k not in cropped_embeddings:
+            cropped_embeddings[k] = list()
+            
+        cropped_embeddings[k].append(v)
+
+clip_dict["frame_numbers"]  =np.concatenate(clip_frames)
+clip_dict["embeds_cropped"] =  {k:np.concatenate(v) for k,v in cropped_embeddings.items()}
+np.savez(crop_embeds_path, **clip_dict)
+

loaders/model_loaders.py

@@ -0,0 +1,127 @@
+import tensorrt as trt
+from cuda.bindings import driver as cuda, runtime as cudart
+import numpy as np
+def cuda_call(call):
+    err, res = call[0], call[1:]
+    check_cuda_err(err)
+    if len(res) == 1:
+        res = res[0]
+    return res
+
+
+def check_cuda_err(err):
+    if isinstance(err, cuda.CUresult):
+        if err != cuda.CUresult.CUDA_SUCCESS:
+            raise RuntimeError("Cuda Error: {}".format(err))
+    if isinstance(err, cudart.cudaError_t):
+        if err != cudart.cudaError_t.cudaSuccess:
+            raise RuntimeError("Cuda Runtime Error: {}".format(err))
+    else:
+        raise RuntimeError("Unknown error type: {}".format(err))
+
+
+
+def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
+    nbytes = host_arr.size * host_arr.itemsize
+    cuda_call(
+        cudart.cudaMemcpy(
+            host_arr, device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost
+        )
+    )
+
+
+def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
+    nbytes = host_arr.size * host_arr.itemsize
+    cuda_call(
+        cudart.cudaMemcpy(
+            device_ptr, host_arr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice
+        )
+    )
+
+
+
+class TensorRTModel:
+    def __init__(self, engine_path):
+        logger = trt.Logger(trt.Logger.ERROR)
+        trt.init_libnvinfer_plugins(logger, namespace="")
+        with open(engine_path, "rb") as f, trt.Runtime(logger) as runtime:
+            assert runtime
+            engine = runtime.deserialize_cuda_engine(f.read())
+        assert engine
+        context = engine.create_execution_context()
+        assert context
+
+        self.engine = engine
+        self.context = context
+        self.logger = logger
+
+    def prepare(self):
+        engine = self.engine
+        inputs = []
+        outputs = []
+        allocations = []
+        for i in range(engine.num_io_tensors):
+            name = engine.get_tensor_name(i)
+            is_input = False
+            if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+                is_input = True
+            dtype = engine.get_tensor_dtype(name)
+            shape = engine.get_tensor_shape(name)
+            if is_input:
+                batch_size = shape[0]
+            size = np.dtype(trt.nptype(dtype)).itemsize
+            for s in shape:
+                size *= s
+            allocation = cuda_call(cudart.cudaMalloc(size))
+            binding = {
+                "index": i,
+                "name": name,
+                "dtype": np.dtype(trt.nptype(dtype)),
+                "shape": list(shape),
+                "allocation": allocation,
+                "size": size,
+            }
+            allocations.append(allocation)
+            if is_input:
+                inputs.append(binding)
+            else:
+                outputs.append(binding)
+
+        assert batch_size > 0
+        assert len(inputs) > 0
+        assert len(outputs) > 0
+        assert len(allocations) > 0
+
+        self.batch_size = batch_size
+        self.inputs = inputs
+        self.outputs = outputs
+        self.allocations = allocations
+
+        assert(len(self.inputs) == 1)
+
+    def allocate_outputs(self):
+        outputs = self.outputs
+        cpu_allocs = dict()
+
+        for output in outputs:
+            dtype = output['dtype']
+            shape = output['shape']
+            name = output['name']
+            output.update({'allocated':     np.empty(shape, dtype=dtype)})
+
+        self.cpu_allocs = cpu_allocs
+
+    def score(self, data_ptr):
+        curr_alloc = self.allocations.copy()
+        curr_alloc[0] = data_ptr
+        self.context.execute_v2(curr_alloc)
+
+        results = dict()
+        for c_out in self.outputs:
+            c_name = c_out['name']
+            c_alloc_idx = c_out['index']
+            c_alloc_array = c_out['allocated']
+            memcpy_device_to_host(c_alloc_array, curr_alloc[c_alloc_idx])
+            results[c_name] = c_alloc_array
+
+        return results

loaders/video_loaders.py

@@ -0,0 +1,25 @@
+import pycuda.driver as cuda_driver
+import PyNvVideoCodec as nvc
+
+gpu_id = 0
+cuda_driver.init()
+cudaDevice = cuda_driver.Device(gpu_id)
+cudaCtx = cudaDevice.retain_primary_context()
+cudaCtx.push()
+cudaStreamNvDec = cuda_driver.Stream()
+
+
+#def create_decoder(enc_file_path, buffer_size = 16):
+    
+decoder = nvc.ThreadedDecoder(
+    enc_file_path = '/home/thebears/Source/task_runners/vision_v3/cuda_objdet_clip/loaders/output.mp4',
+    buffer_size=32,
+    gpu_id=gpu_id,
+    cuda_context=cudaCtx.handle,
+    cuda_stream=cudaStreamNvDec.handle,
+    use_device_memory=True,
+    output_color_type=nvc.OutputColorType.RGBP,
+)
+ #   return  decoder
+
+

models/.gitattributes

@@ -0,0 +1,5 @@
+dfine_large_8.6.engine filter=lfs diff=lfs merge=lfs -text
+dfine_large_8.9.engine filter=lfs diff=lfs merge=lfs -text
+dfine_large.species_keep filter=lfs diff=lfs merge=lfs -text
+ViT-SO400M-16-SigLIP2-512_visual_fp16_8.6.engine filter=lfs diff=lfs merge=lfs -text
+ViT-SO400M-16-SigLIP2-512_visual_fp16_8.9.engine filter=lfs diff=lfs merge=lfs -text

transforms/model_transforms.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn.functional as F
+from torchvision.transforms import v2
+
+
+class ResizeAndPadBatch:
+    """
+    Vectorized resize + zero pad for batched tensors.
+
+    Args:
+        imgs (Tensor): shape [B, C, H, W], pixel values in [0,1] or [0,255].
+        target_size (tuple): (target_h, target_w) final image shape.
+
+    Returns:
+        Tensor: shape [B, C, target_h, target_w]
+    """
+
+    def __init__(self, target_size=(640, 640), fill_value=0):
+        self.target_size = target_size
+        self.fill_value = fill_value
+        self.transform_stats = None
+
+    def __call__(self, imgs: torch.Tensor):
+
+        b, c, h, w = imgs.shape
+        target_h, target_w = self.target_size
+
+        # Compute scaling factors: preserve aspect ratio by same scale factor per image
+        scale_factors = float(
+            torch.minimum(torch.tensor(target_h / h), torch.tensor(target_w / w))
+        )
+
+        # New intermediate size
+        new_h = int(h * scale_factors)
+        new_w = int(w * scale_factors)
+
+        # Resize with bilinear interpolation
+        resized = F.interpolate(
+            imgs, size=(new_h, new_w), mode="bilinear", align_corners=False
+        )
+
+        # Calculate padding amounts (left, right, top, bottom)
+        pad_h = target_h - new_h
+        pad_w = target_w - new_w
+        pad_top = pad_h // 2
+        pad_bottom = pad_h - pad_top
+        pad_left = pad_w // 2
+        pad_right = pad_w - pad_left
+
+        # Apply padding (pad order in F.pad is [left, right, top, bottom])
+        padded = F.pad(
+            resized,
+            (pad_left, pad_right, pad_top, pad_bottom),
+            mode="constant",
+            value=self.fill_value,
+        )
+
+        self.transform_stats = {
+            "scale_factor": scale_factors,
+            'hw': [h,w],
+            "new_hw": [new_h, new_w],
+            "pad_TBLR": [pad_top,pad_bottom, pad_left, pad_right]
+        }
+
+        return padded
+
+
+clip_transforms = v2.Compose(
+    [
+        v2.Resize(size=(512, 512)),
+        v2.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+    ]
+)
+
+det_transforms = v2.Compose([ResizeAndPadBatch(target_size=(640, 640))])