|
| 1 | +from vlmeval.dataset.video_base import VideoBaseDataset |
| 2 | +from huggingface_hub import hf_hub_download |
| 3 | +from datasets import load_dataset |
| 4 | +from ..smp.file import load |
| 5 | +from PIL import Image |
| 6 | +import numpy as np |
| 7 | +import portalocker |
| 8 | +import zipfile |
| 9 | +import ast |
| 10 | +import os |
| 11 | + |
| 12 | +NQ_QUESTION_TYPES = [ |
| 13 | + "object_counting_single", |
| 14 | + "object_counting_multiple", |
| 15 | + "object_abs_distance", |
| 16 | + "object_size_estimation", |
| 17 | + "room_size_estimation_single", |
| 18 | + "room_size_estimation_multiple" |
| 19 | +] |
| 20 | + |
| 21 | + |
| 22 | +MCQ_QUESTION_TYPES = [ |
| 23 | + "object_rel_direction_forward_easy", |
| 24 | + "object_rel_direction_backward_easy", |
| 25 | + "object_rel_direction_forward_hard", |
| 26 | + "object_rel_direction_backward_hard", |
| 27 | + "object_rel_distance_closest", |
| 28 | + "object_rel_distance_farthest", |
| 29 | + "route_planning" |
| 30 | +] |
| 31 | + |
| 32 | +PROMPT_PREFIX = "These are frames of a video." |
| 33 | +REPO_ID = "3dlg-hcvc/ReVSI" |
| 34 | +VIDEO_ROOT_DIR = "revsi_videos" |
| 35 | +EXTRACT_SENTINEL = ".extracted" |
| 36 | + |
| 37 | + |
| 38 | +def _safe_extract_zip(zip_path, target_dir): |
| 39 | + target_dir_abs = os.path.abspath(target_dir) |
| 40 | + with zipfile.ZipFile(zip_path, "r") as zf: |
| 41 | + for info in zf.infolist(): |
| 42 | + rel_path = os.path.normpath(info.filename).lstrip("/\\") |
| 43 | + dst_path = os.path.abspath(os.path.join(target_dir, rel_path)) |
| 44 | + if not dst_path.startswith(target_dir_abs + os.sep): |
| 45 | + raise RuntimeError(f"Unsafe path in zip: {info.filename}") |
| 46 | + if info.is_dir(): |
| 47 | + os.makedirs(dst_path, exist_ok=True) |
| 48 | + continue |
| 49 | + os.makedirs(os.path.dirname(dst_path), exist_ok=True) |
| 50 | + with zf.open(info, "r") as src, open(dst_path, "wb") as dst: |
| 51 | + dst.write(src.read()) |
| 52 | + |
| 53 | + |
| 54 | +def _write_sentinel(path): |
| 55 | + tmp_path = path + ".tmp" |
| 56 | + with open(tmp_path, "w", encoding="utf-8") as f: |
| 57 | + f.write("done") |
| 58 | + os.replace(tmp_path, path) |
| 59 | + |
| 60 | + |
| 61 | +def _serialize_options(options): |
| 62 | + if isinstance(options, np.ndarray): |
| 63 | + options = options.tolist() |
| 64 | + if isinstance(options, (list, tuple)): |
| 65 | + return repr([str(option) for option in options]) |
| 66 | + return options |
| 67 | + |
| 68 | + |
| 69 | +def _parse_options(options): |
| 70 | + if isinstance(options, np.ndarray): |
| 71 | + return [str(option) for option in options.tolist()] |
| 72 | + if isinstance(options, (list, tuple)): |
| 73 | + return [str(option) for option in options] |
| 74 | + if isinstance(options, str): |
| 75 | + return [str(option) for option in ast.literal_eval(options)] |
| 76 | + return [str(options)] |
| 77 | + |
| 78 | + |
| 79 | +def _mean_relative_accuracy(pred, target, start, end, interval): |
| 80 | + num_pts = (end - start) / interval + 2 |
| 81 | + conf_intervs = np.linspace(start, end, int(num_pts)) |
| 82 | + accuracy = (abs(pred - target) / target) <= (1 - conf_intervs) |
| 83 | + return accuracy.mean() |
| 84 | + |
| 85 | + |
| 86 | +def _pop_mean(output, metric_name, metric_keys): |
| 87 | + values = [output.pop(key) for key in metric_keys if key in output] |
| 88 | + if values: |
| 89 | + output[metric_name] = np.mean(values) |
| 90 | + |
| 91 | + |
| 92 | +class ReVSI(VideoBaseDataset): |
| 93 | + |
| 94 | + TYPE = 'Video-VQA' |
| 95 | + |
| 96 | + def __init__(self, dataset='ReVSI', pack=False, nframe=None, **kwargs): |
| 97 | + if nframe in [None, "all", "all_frame"]: |
| 98 | + self.frame_subset = "all_frame" |
| 99 | + nframe = 128 |
| 100 | + else: |
| 101 | + nframe = int(nframe) |
| 102 | + self.frame_subset = f"{nframe}_frame" |
| 103 | + super().__init__(dataset=dataset, pack=pack, nframe=nframe, fps=-1) |
| 104 | + |
| 105 | + @classmethod |
| 106 | + def supported_datasets(cls): |
| 107 | + return ['ReVSI'] |
| 108 | + |
| 109 | + def prepare_dataset(self, dataset): |
| 110 | + subset = self.frame_subset |
| 111 | + dataset_table = load_dataset(REPO_ID, subset, split="test") |
| 112 | + dataset_table = dataset_table.add_column('video', [f"{x['scene_id']}.mp4" for x in dataset_table]) |
| 113 | + df = dataset_table.to_pandas() |
| 114 | + if 'options' in df: |
| 115 | + df['options'] = df['options'].apply(_serialize_options) |
| 116 | + video_zip_path = hf_hub_download(repo_id=REPO_ID, filename="video.zip", repo_type="dataset") |
| 117 | + dataset_path = os.path.dirname(video_zip_path) |
| 118 | + video_root = os.path.join(dataset_path, VIDEO_ROOT_DIR) |
| 119 | + os.makedirs(video_root, exist_ok=True) |
| 120 | + sentinel_path = os.path.join(video_root, EXTRACT_SENTINEL) |
| 121 | + expected_video_path = os.path.join(video_root, subset, df["video"].iloc[0]) |
| 122 | + |
| 123 | + def videos_ready(): |
| 124 | + return os.path.exists(sentinel_path) and os.path.isfile(expected_video_path) |
| 125 | + |
| 126 | + if not videos_ready(): |
| 127 | + lock_path = os.path.join(video_root, ".extract.lock") |
| 128 | + with portalocker.Lock(lock_path, "w", timeout=300): |
| 129 | + if not videos_ready(): |
| 130 | + _safe_extract_zip(video_zip_path, video_root) |
| 131 | + _write_sentinel(sentinel_path) |
| 132 | + |
| 133 | + tsv_file_path = os.path.join(dataset_path, f"{subset}.tsv") |
| 134 | + df.to_csv(tsv_file_path, sep="\t", index=False) |
| 135 | + return dict(root=video_root, data_file=tsv_file_path) |
| 136 | + |
| 137 | + def video_path(self, video): |
| 138 | + return os.path.join(self.data_root, self.frame_subset, video) |
| 139 | + |
| 140 | + def save_video_frames(self, video): |
| 141 | + import decord |
| 142 | + |
| 143 | + vid_path = self.video_path(video) |
| 144 | + frame_key = os.path.join(self.frame_subset, os.path.splitext(video)[0]) |
| 145 | + vid = decord.VideoReader(vid_path) |
| 146 | + video_fps = vid.get_avg_fps() |
| 147 | + if self.nframe > 0 and self.fps < 0: |
| 148 | + step_size = len(vid) / (self.nframe + 1) |
| 149 | + indices = [int(i * step_size) for i in range(1, self.nframe + 1)] |
| 150 | + frame_paths = self.frame_paths(frame_key) |
| 151 | + lock_name = f"{os.path.splitext(video)[0]}.{self.nframe}frame.lock" |
| 152 | + elif self.fps > 0: |
| 153 | + total_duration = len(vid) / video_fps |
| 154 | + required_frames = max(int(total_duration * self.fps), 1) |
| 155 | + step_size = video_fps / self.fps |
| 156 | + indices = [min(int(i * step_size), len(vid) - 1) for i in range(required_frames)] |
| 157 | + frame_paths = self.frame_paths_fps(frame_key, len(indices)) |
| 158 | + lock_name = f"{os.path.splitext(video)[0]}.{self.fps}fps.lock" |
| 159 | + else: |
| 160 | + raise ValueError('ReVSI requires either nframe > 0 or fps > 0 to extract frames') |
| 161 | + |
| 162 | + if np.all([os.path.exists(p) for p in frame_paths]): |
| 163 | + return frame_paths |
| 164 | + |
| 165 | + lock_dir = os.path.join(self.frame_root, self.frame_subset) |
| 166 | + os.makedirs(lock_dir, exist_ok=True) |
| 167 | + lock_path = os.path.join(lock_dir, lock_name) |
| 168 | + with portalocker.Lock(lock_path, "w", timeout=30): |
| 169 | + if np.all([os.path.exists(p) for p in frame_paths]): |
| 170 | + return frame_paths |
| 171 | + images = [Image.fromarray(vid[i].asnumpy()) for i in indices] |
| 172 | + for image, path in zip(images, frame_paths): |
| 173 | + if not os.path.exists(path): |
| 174 | + image.save(path) |
| 175 | + return frame_paths |
| 176 | + |
| 177 | + def build_prompt(self, idx, video_llm): |
| 178 | + line = self.data.iloc[idx] |
| 179 | + question_type = line["question_type"] |
| 180 | + question = line["question"] |
| 181 | + if question_type in NQ_QUESTION_TYPES: |
| 182 | + post_prompt = "Answer the question using a single integer or decimal number." |
| 183 | + full_prompt = "\n".join([PROMPT_PREFIX, question, post_prompt]).strip() |
| 184 | + elif question_type in MCQ_QUESTION_TYPES: |
| 185 | + options = _parse_options(line["options"]) |
| 186 | + options_str = "Options:\n" + "\n".join(options) |
| 187 | + post_prompt = "Answer with the option's letter from the given choices directly." |
| 188 | + full_prompt = "\n".join([PROMPT_PREFIX, question, options_str, post_prompt]).strip() |
| 189 | + message = [] |
| 190 | + if video_llm: |
| 191 | + message.append(dict(type='video', value=self.video_path(line["video"]))) |
| 192 | + else: |
| 193 | + for frame_path in self.save_video_frames(line["video"]): |
| 194 | + message.append(dict(type='image', value=frame_path)) |
| 195 | + message.append(dict(type='text', value=full_prompt)) |
| 196 | + return message |
| 197 | + |
| 198 | + def evaluate(self, eval_file, **judge_kwargs): |
| 199 | + df = load(eval_file) |
| 200 | + for i, row in df.iterrows(): |
| 201 | + pred_answer = str(row["prediction"]).strip().split(" ")[0].rstrip(".").strip() |
| 202 | + gt_answer = str(row["ground_truth"]) |
| 203 | + if row["question_type"] in MCQ_QUESTION_TYPES: |
| 204 | + accuracy = 1.0 if pred_answer.lower() == gt_answer.lower() else 0.0 |
| 205 | + elif row["question_type"] in NQ_QUESTION_TYPES: |
| 206 | + try: |
| 207 | + accuracy = _mean_relative_accuracy( |
| 208 | + float(pred_answer), float(gt_answer), 0.5, 0.95, 0.05 |
| 209 | + ) |
| 210 | + except (TypeError, ValueError, ZeroDivisionError): |
| 211 | + accuracy = 0.0 |
| 212 | + df.at[i, "accuracy"] = accuracy |
| 213 | + |
| 214 | + output = {} |
| 215 | + for question_type, per_question_type in df.groupby("question_type"): |
| 216 | + output[f"{question_type}_accuracy"] = per_question_type["accuracy"].mean() |
| 217 | + |
| 218 | + _pop_mean(output, "object_rel_direction_accuracy", [ |
| 219 | + "object_rel_direction_forward_easy_accuracy", |
| 220 | + "object_rel_direction_backward_easy_accuracy", |
| 221 | + "object_rel_direction_forward_hard_accuracy", |
| 222 | + "object_rel_direction_backward_hard_accuracy", |
| 223 | + ]) |
| 224 | + _pop_mean(output, "object_rel_distance_accuracy", [ |
| 225 | + "object_rel_distance_closest_accuracy", |
| 226 | + "object_rel_distance_farthest_accuracy", |
| 227 | + ]) |
| 228 | + _pop_mean(output, "object_counting_accuracy", [ |
| 229 | + "object_counting_single_accuracy", |
| 230 | + "object_counting_multiple_accuracy", |
| 231 | + ]) |
| 232 | + _pop_mean(output, "room_size_estimation_accuracy", [ |
| 233 | + "room_size_estimation_single_accuracy", |
| 234 | + "room_size_estimation_multiple_accuracy", |
| 235 | + ]) |
| 236 | + output["overall_accuracy"] = sum(output.values()) / len(output) if output else 0.0 |
| 237 | + return output |
0 commit comments