commit f18db3f45bc96072666be1545da9b0539f7539d2 Author: Ishan S. Patel Date: Wed Jul 8 10:57:49 2026 -0400 adding models diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..e69de29 diff --git a/__pycache__/infer.cpython-312.pyc b/__pycache__/infer.cpython-312.pyc new file mode 100644 index 0000000..faabda7 Binary files /dev/null and b/__pycache__/infer.cpython-312.pyc differ diff --git a/infer.py b/infer.py new file mode 100644 index 0000000..19eb00f --- /dev/null +++ b/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 + + diff --git a/infer_sub_clip.py b/infer_sub_clip.py new file mode 100644 index 0000000..0eb5ee7 --- /dev/null +++ b/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) + diff --git a/loaders/__pycache__/model_loaders.cpython-312.pyc b/loaders/__pycache__/model_loaders.cpython-312.pyc new file mode 100644 index 0000000..6d30057 Binary files /dev/null and b/loaders/__pycache__/model_loaders.cpython-312.pyc differ diff --git a/loaders/__pycache__/video_loaders.cpython-312.pyc b/loaders/__pycache__/video_loaders.cpython-312.pyc new file mode 100644 index 0000000..9875784 Binary files /dev/null and b/loaders/__pycache__/video_loaders.cpython-312.pyc differ diff --git a/loaders/__pycache__/video_loaders.cpython-313.pyc b/loaders/__pycache__/video_loaders.cpython-313.pyc new file mode 100644 index 0000000..240c666 Binary files /dev/null and b/loaders/__pycache__/video_loaders.cpython-313.pyc differ diff --git a/loaders/model_loaders.py b/loaders/model_loaders.py new file mode 100644 index 0000000..42e7a94 --- /dev/null +++ b/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 diff --git a/loaders/output.mp4 b/loaders/output.mp4 new file mode 100644 index 0000000..c7a87df Binary files /dev/null and b/loaders/output.mp4 differ diff --git a/loaders/video_loaders.py b/loaders/video_loaders.py new file mode 100644 index 0000000..0d807f8 --- /dev/null +++ b/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 + + diff --git a/models/.gitattributes b/models/.gitattributes new file mode 100644 index 0000000..f6e5988 --- /dev/null +++ b/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 diff --git a/models/ViT-SO400M-16-SigLIP2-512_visual_fp16_8.6.engine b/models/ViT-SO400M-16-SigLIP2-512_visual_fp16_8.6.engine new file mode 100644 index 0000000..74abbff --- /dev/null +++ b/models/ViT-SO400M-16-SigLIP2-512_visual_fp16_8.6.engine @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:705c98eec0664649ae0dc73ff913e5e6e9c7b623dcc74ef12841296dd146e88a +size 865881748 diff --git a/models/ViT-SO400M-16-SigLIP2-512_visual_fp16_8.9.engine b/models/ViT-SO400M-16-SigLIP2-512_visual_fp16_8.9.engine new file mode 100644 index 0000000..b6ff67d --- /dev/null +++ b/models/ViT-SO400M-16-SigLIP2-512_visual_fp16_8.9.engine @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5c53aa5217e4b0db270af73da786c38878b34abe2bff7d4860bb21f67caa8419 +size 861264132 diff --git a/models/dfine_large.species_keep b/models/dfine_large.species_keep new file mode 100644 index 0000000..a11d9b7 --- /dev/null +++ b/models/dfine_large.species_keep @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:795a1ccbc813b98dc29c446d1bf50badcde6dea1ef0dd6bae95061e7f68b4fbb +size 3736 diff --git a/models/dfine_large_8.6.engine b/models/dfine_large_8.6.engine new file mode 100644 index 0000000..3454457 --- /dev/null +++ b/models/dfine_large_8.6.engine @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0d08b296dea3b1445b708c372ecf34889585b93b652f97f794caacd95a1008c3 +size 74004364 diff --git a/models/dfine_large_8.9.engine b/models/dfine_large_8.9.engine new file mode 100644 index 0000000..ca5503b --- /dev/null +++ b/models/dfine_large_8.9.engine @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ba3010a2b787df00a53f80d3d55cda60017b25d65d8c6772875843665f9e038f +size 72831484 diff --git a/transforms/__pycache__/model_transforms.cpython-312.pyc b/transforms/__pycache__/model_transforms.cpython-312.pyc new file mode 100644 index 0000000..0178658 Binary files /dev/null and b/transforms/__pycache__/model_transforms.cpython-312.pyc differ diff --git a/transforms/model_transforms.py b/transforms/model_transforms.py new file mode 100644 index 0000000..18747ba --- /dev/null +++ b/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))])