Buckets:
| #!/usr/bin/env python3 | |
| """ | |
| Randomly sample N videos and generate cropped reference images from bbox annotations. | |
| Usage: | |
| python generate_reference_images.py --type face | |
| python generate_reference_images.py --type all --num 3000 --output ref_imgs | |
| """ | |
| import argparse | |
| import json | |
| import random | |
| import sys | |
| from pathlib import Path | |
| import cv2 | |
| BASE_DIR = Path(__file__).parent | |
| VIDEOS_DIR = BASE_DIR / "videos" / "videos" | |
| BBOXES_DIR = BASE_DIR / "bboxes" / "refine_bbox_jsons" | |
| ALL_TYPES = ["face", "head", "person"] | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description="Generate reference images from bbox annotations") | |
| parser.add_argument( | |
| "--type", | |
| choices=["face", "head", "person", "all"], | |
| default="face", | |
| help="Bbox type to crop: face | head | person | all (default: face)", | |
| ) | |
| parser.add_argument("--num", type=int, default=3000, help="Number of videos to sample (default: 3000)") | |
| parser.add_argument("--output", type=str, default="reference_images", help="Output directory (default: reference_images)") | |
| parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)") | |
| parser.add_argument( | |
| "--masks", | |
| type=str, | |
| default=None, | |
| help="Directory of mask images (named {stem}.png/.jpg). Non-masked area in the crop will be set to white.", | |
| ) | |
| return parser.parse_args() | |
| def find_matched_pairs(): | |
| video_by_stem = {p.stem: p for p in VIDEOS_DIR.glob("*.mp4")} | |
| json_by_stem = {p.stem: p for p in BBOXES_DIR.glob("*.json")} | |
| matched = [(video_by_stem[s], json_by_stem[s]) for s in video_by_stem if s in json_by_stem] | |
| return matched | |
| def frames_with_all_types(bbox_data, types): | |
| """Return frame keys where every requested type has at least one detection.""" | |
| valid = [] | |
| for frame_key, frame_data in bbox_data.items(): | |
| if all(frame_data.get(t) for t in types): | |
| valid.append(frame_key) | |
| return valid | |
| def read_frame(video_path, frame_idx): | |
| cap = cv2.VideoCapture(str(video_path)) | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, int(frame_idx)) | |
| ret, frame = cap.read() | |
| cap.release() | |
| return frame if ret else None | |
| def crop_box(frame, box): | |
| h, w = frame.shape[:2] | |
| x1 = max(0, int(box["x1"])) | |
| y1 = max(0, int(box["y1"])) | |
| x2 = min(w, int(box["x2"])) | |
| y2 = min(h, int(box["y2"])) | |
| return frame[y1:y2, x1:x2] | |
| def resolve_masks_dir(masks_arg): | |
| """Accept either the parent masks dir or the track_masks_data subdir directly.""" | |
| p = Path(masks_arg) | |
| candidate = p / "track_masks_data" | |
| return candidate if candidate.exists() else p | |
| def load_mask_info(masks_dir, stem): | |
| """Load mask metadata for a video stem. | |
| Returns (corresponding_data, valid_frames_data, stem_dir) or None if unavailable. | |
| corresponding_data: {person_id: {type: object_id}} | |
| valid_frames_data: {person_id: {type: [frame_indices]}} | |
| """ | |
| stem_dir = masks_dir / stem | |
| if not stem_dir.exists(): | |
| return None | |
| try: | |
| with open(stem_dir / "corresponding_data.json") as f: | |
| corresponding = json.load(f) | |
| with open(stem_dir / "valid_frame.json") as f: | |
| valid_frames_data = json.load(f) | |
| return corresponding, valid_frames_data, stem_dir | |
| except Exception: | |
| return None | |
| def mask_valid_frame_keys(mask_info, types): | |
| """Return frame keys (strings) valid for all requested types in the mask data.""" | |
| corresponding, valid_frames_data, _ = mask_info | |
| person_id = next(iter(corresponding)) | |
| vf = valid_frames_data.get(person_id, {}) | |
| valid = None | |
| for t in types: | |
| frames = {str(idx) for idx in vf.get(t, [])} | |
| valid = frames if valid is None else valid & frames | |
| return valid or set() | |
| def read_type_mask(mask_info, frame_key, t): | |
| """Read the uint16 mask for a specific type and frame. Returns array or None.""" | |
| corresponding, _, stem_dir = mask_info | |
| person_id = next(iter(corresponding)) | |
| obj_id = corresponding[person_id].get(t) | |
| if obj_id is None: | |
| return None | |
| mask_path = stem_dir / "tracking_mask_results" / str(obj_id) / f"annotated_frame_{int(frame_key):05d}.png" | |
| if not mask_path.exists(): | |
| return None | |
| return cv2.imread(str(mask_path), cv2.IMREAD_UNCHANGED) | |
| def apply_mask_to_crop(crop, mask, box): | |
| """Crop mask to bbox region and set zero-mask pixels in crop to white.""" | |
| h, w = mask.shape[:2] | |
| x1 = max(0, int(box["x1"])) | |
| y1 = max(0, int(box["y1"])) | |
| x2 = min(w, int(box["x2"])) | |
| y2 = min(h, int(box["y2"])) | |
| mask_crop = mask[y1:y2, x1:x2] | |
| if mask_crop.shape[:2] != crop.shape[:2]: | |
| mask_crop = cv2.resize(mask_crop, (crop.shape[1], crop.shape[0]), interpolation=cv2.INTER_NEAREST) | |
| result = crop.copy() | |
| result[mask_crop == 0] = 255 | |
| return result | |
| def best_detection(detections): | |
| """Return detection with highest confidence.""" | |
| return max(detections, key=lambda d: d["confidence"]) | |
| def main(): | |
| args = parse_args() | |
| random.seed(args.seed) | |
| types = ALL_TYPES if args.type == "all" else [args.type] | |
| output_dir = Path(args.output) | |
| masks_dir = resolve_masks_dir(args.masks) if args.masks else None | |
| for t in types: | |
| (output_dir / t).mkdir(parents=True, exist_ok=True) | |
| print(f"Scanning {VIDEOS_DIR} for matched pairs...") | |
| matched = find_matched_pairs() | |
| print(f"Found {len(matched)} video/bbox pairs") | |
| n = min(args.num, len(matched)) | |
| if n < args.num: | |
| print(f"Warning: only {len(matched)} pairs available, using all of them") | |
| sampled = random.sample(matched, n) | |
| print(f"Processing {n} videos [type={args.type}] -> {output_dir}/") | |
| success, skipped = 0, 0 | |
| for i, (video_path, json_path) in enumerate(sampled): | |
| if (i + 1) % 200 == 0: | |
| print(f" [{i+1}/{n}] success={success} skipped={skipped}") | |
| try: | |
| with open(json_path) as f: | |
| bbox_data = json.load(f) | |
| except Exception as e: | |
| print(f" skip {json_path.name}: {e}", file=sys.stderr) | |
| skipped += 1 | |
| continue | |
| stem = video_path.stem | |
| valid_frames = set(frames_with_all_types(bbox_data, types)) | |
| mask_info = None | |
| if masks_dir is not None: | |
| mask_info = load_mask_info(masks_dir, stem) | |
| if mask_info is not None: | |
| valid_frames &= mask_valid_frame_keys(mask_info, types) | |
| if not valid_frames: | |
| skipped += 1 | |
| continue | |
| frame_key = random.choice(list(valid_frames)) | |
| frame = read_frame(video_path, frame_key) | |
| if frame is None: | |
| skipped += 1 | |
| continue | |
| frame_data = bbox_data[frame_key] | |
| wrote_all = True | |
| for t in types: | |
| det = best_detection(frame_data[t]) | |
| crop = crop_box(frame, det["box"]) | |
| if crop.size == 0: | |
| wrote_all = False | |
| break | |
| if mask_info is not None: | |
| type_mask = read_type_mask(mask_info, frame_key, t) | |
| if type_mask is not None: | |
| crop = apply_mask_to_crop(crop, type_mask, det["box"]) | |
| cv2.imwrite(str(output_dir / t / f"{stem}.jpg"), crop) | |
| if wrote_all: | |
| success += 1 | |
| else: | |
| skipped += 1 | |
| print(f"\nDone: {success} saved, {skipped} skipped") | |
| for t in types: | |
| saved = list((output_dir / t).glob("*.jpg")) | |
| print(f" {output_dir / t}: {len(saved)} images") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 7.59 kB
- Xet hash:
- fcd42ef0074a74e30e96c6c715b149d3330fafa15cf0ea5745e5584c4df2b1ad
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.