File size: 5,521 Bytes
0c51b93 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
from operator import attrgetter
from llava.model.builder import load_pretrained_model
from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, IGNORE_INDEX
from llava.conversation import conv_templates, SeparatorStyle
import torch
import cv2
import numpy as np
from PIL import Image
import requests
import copy
import warnings
from decord import VideoReader, cpu
import re
from tqdm import tqdm
warnings.filterwarnings("ignore")
import torch.multiprocessing as mp
mp.set_start_method('spawn', force=True)
pretrained = "your_llava-ov_checkpoint"
model_name = "llava_qwen"
llava_model_args = {
"multimodal": True,
}
def load_model(device):
"""Initialize model on a specific GPU."""
tokenizer, model, image_processor, max_length = load_pretrained_model(
pretrained, None, model_name, device_map={"": device}, attn_implementation="sdpa", **llava_model_args
)
model.eval()
return tokenizer, model, image_processor
def load_video(video_path, max_frames_num):
"""Function to extract frames from video"""
if isinstance(video_path, str):
vr = VideoReader(video_path, ctx=cpu(0))
else:
vr = VideoReader(video_path[0], ctx=cpu(0))
total_frame_num = len(vr)
uniform_sampled_frames = np.linspace(0, total_frame_num - 1, max_frames_num, dtype=int)
frame_idx = uniform_sampled_frames.tolist()
spare_frames = vr.get_batch(frame_idx).asnumpy()
return spare_frames
def split_list(lst, num_splits):
"""Split a list into evenly sized chunks"""
chunk_size = len(lst) // num_splits
remainder = len(lst) % num_splits
chunks = []
start = 0
for i in range(num_splits):
end = start + chunk_size + (1 if i < remainder else 0)
chunks.append(lst[start:end])
start = end
return chunks
def process_chunk(gpu_id, chunk, results):
"""Process a chunk of data on a specific GPU."""
device = f"cuda:{gpu_id}"
tokenizer, model, image_processor = load_model(device)
model = model.to(device)
res_l = []
for sample in tqdm(chunk, desc=f"GPU {gpu_id}"):
try:
# Load and process video
video_path = sample[0]
video_frames = load_video(video_path, 32)
frames = image_processor.preprocess(video_frames, return_tensors="pt")["pixel_values"].half().to(device)
# Prepare conversation input
conv_template = "qwen_1_5"
question = f"{DEFAULT_IMAGE_TOKEN}\nAnalyze the video content and evaluate the performance for the following 20 skills. Assign a ranking to each skill on a scale of 0 to 4, where 0 represents no impact or poor performance and 4 represents excellent performance. Follow the format: `Skill Name: Ranking` (e.g., `Free Throw: 4`). The skills to evaluate are: Free Throw, 2-PTs, 3-PTs, Contested-shots, Overall Shooting, Rebounds, Defensive Rebounds, Offensive Rebounds, Steals, Turnovers, Points-allowed, Defensive Consistency, Assists, Passing Accuracy, Fouls, Contribution, Offensive Consistency, Teamwork, Impact, Efficiency"
conv = copy.deepcopy(conv_templates[conv_template])
conv.append_message(conv.roles[0], question)
conv.append_message(conv.roles[1], None)
prompt_question = conv.get_prompt()
input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(device)
image_sizes = [frame.size for frame in video_frames]
# Generate response
cont = model.generate(
input_ids,
images=[frames],
image_sizes=image_sizes,
do_sample=False,
temperature=0,
max_new_tokens=4096,
modalities=["video"],
)
text_outputs = tokenizer.batch_decode(cont, skip_special_tokens=True)
src_txt = text_outputs[0].replace('\n', ',')
src_num = re.findall(r':\s*(\d+)', src_txt)
src_num = [int(num) for num in src_num]
if len(src_num) != 20:
print(f'output skill num mismatch! {len(src_num)}')
tgt_num = [int(x) for x in sample[1:]]
matches = sum(1 for x, y in zip(src_num, tgt_num) if x == y) / len(tgt_num)
res_l.append(matches)
except Exception as e:
print(f"Error on GPU {gpu_id}: {e}")
continue
results[gpu_id] = res_l
if __name__ == "__main__":
test_path = 'your_BASKET_test_path'
test_list = []
with open(test_path, 'r') as f:
for line in f.readlines()[1:-2]:
test_list.append(line.strip().split(','))
num_gpus = torch.cuda.device_count()
if num_gpus == 0:
raise RuntimeError("No GPUs available")
test_chunks = split_list(test_list, num_gpus)
manager = mp.Manager()
results = manager.dict()
processes = []
for gpu_id, chunk in enumerate(test_chunks):
p = mp.Process(target=process_chunk, args=(gpu_id, chunk, results))
processes.append(p)
p.start()
for p in processes:
p.join()
# Aggregate results
all_results = []
for res in results.values():
all_results.extend(res)
average = sum(all_results) / len(all_results)
print(f'Average accuracy out of 20: {average}')
|