110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
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)
|
|
|