adding models
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user