File size: 3,429 Bytes
f53b39e |
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 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import random
from dataclasses import dataclass
from typing import List
from training.dataset.vos_segment_loader import LazySegments
MAX_RETRIES = 1000
@dataclass
class SampledFramesAndObjects:
frames: List[int]
object_ids: List[int]
class VOSSampler:
def __init__(self, sort_frames=True):
# frames are ordered by frame id when sort_frames is True
self.sort_frames = sort_frames
def sample(self, video):
raise NotImplementedError()
class RandomUniformSampler(VOSSampler):
def __init__(
self,
num_frames,
max_num_objects,
reverse_time_prob=0.0,
):
self.num_frames = num_frames
self.max_num_objects = max_num_objects
self.reverse_time_prob = reverse_time_prob
def sample(self, video, segment_loader, epoch=None):
for retry in range(MAX_RETRIES):
if len(video.frames) < self.num_frames:
raise Exception(
f"Cannot sample {self.num_frames} frames from video {video.video_name} as it only has {len(video.frames)} annotated frames."
)
start = random.randrange(0, len(video.frames) - self.num_frames + 1)
frames = [video.frames[start + step] for step in range(self.num_frames)]
if random.uniform(0, 1) < self.reverse_time_prob:
# Reverse time
frames = frames[::-1]
# Get first frame object ids
visible_object_ids = []
loaded_segms = segment_loader.load(frames[0].frame_idx)
if isinstance(loaded_segms, LazySegments):
# LazySegments for SA1BRawDataset
visible_object_ids = list(loaded_segms.keys())
else:
for object_id, segment in segment_loader.load(
frames[0].frame_idx
).items():
if segment.sum():
visible_object_ids.append(object_id)
# First frame needs to have at least a target to track
if len(visible_object_ids) > 0:
break
if retry >= MAX_RETRIES - 1:
raise Exception("No visible objects")
object_ids = random.sample(
visible_object_ids,
min(len(visible_object_ids), self.max_num_objects),
)
return SampledFramesAndObjects(frames=frames, object_ids=object_ids)
class EvalSampler(VOSSampler):
"""
VOS Sampler for evaluation: sampling all the frames and all the objects in a video
"""
def __init__(
self,
):
super().__init__()
def sample(self, video, segment_loader, epoch=None):
"""
Sampling all the frames and all the objects
"""
if self.sort_frames:
# ordered by frame id
frames = sorted(video.frames, key=lambda x: x.frame_idx)
else:
# use the original order
frames = video.frames
object_ids = segment_loader.load(frames[0].frame_idx).keys()
if len(object_ids) == 0:
raise Exception("First frame of the video has no objects")
return SampledFramesAndObjects(frames=frames, object_ids=object_ids)
|