306 lines
10 KiB
Python
306 lines
10 KiB
Python
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
|
|
|
|
|