(video_file, preprocess, model, batch_size, device)
| 82 | |
| 83 | |
| 84 | def encode_video(video_file, preprocess, model, batch_size, device): |
| 85 | cv_start_time = time.perf_counter() |
| 86 | cap = cv2.VideoCapture(video_file) |
| 87 | frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| 88 | images = [] |
| 89 | count = 0 |
| 90 | ret = True |
| 91 | |
| 92 | while (count < frameCount and ret): |
| 93 | ret, frame = cap.read() |
| 94 | if not ret: # if file is empty break loop |
| 95 | break |
| 96 | frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 97 | images.append(preprocess(Image.fromarray(frame_rgb).convert("RGB"))) |
| 98 | count += 1 |
| 99 | |
| 100 | cv_end_time = time.perf_counter() |
| 101 | time_diff = cv_end_time - cv_start_time |
| 102 | # print(f"cv done in {time_diff:.2f} seconds") |
| 103 | |
| 104 | image_embed_start_time = time.perf_counter() |
| 105 | image_input = torch.tensor(np.stack(images)).to(device) |
| 106 | image_features_list = [] |
| 107 | # bs = 256 |
| 108 | with torch.no_grad(): |
| 109 | n_inter = math.ceil(len(image_input) / batch_size) |
| 110 | for i in range(n_inter): |
| 111 | image_features = model.encode_image(image_input[i * batch_size: (i + 1) * batch_size]).float() |
| 112 | image_features_list.append(image_features) |
| 113 | image_features = torch.cat(image_features_list, dim=0) |
| 114 | image_features /= image_features.norm(dim=-1, keepdim=True) |
| 115 | cap.release() |
| 116 | |
| 117 | vid_feature = normalize_matrix(torch.mean(image_features, dim=0, keepdim=True)).squeeze() |
| 118 | |
| 119 | image_embed_end_time = time.perf_counter() |
| 120 | time_diff = image_embed_end_time - image_embed_start_time |
| 121 | # print(f"image embed done in {time_diff:.2f} seconds") |
| 122 | |
| 123 | return image_features, vid_feature |
| 124 | |
| 125 | |
| 126 | def encode_text(vid_caps, model, tokenizer, idf_dict, device): |
no test coverage detected