ssn_dataset.py

import torch.utils.data as data

import os
import os.path
from numpy.random import randint
from ops.io import load_proposal_file
from transforms import *
from ops.utils import temporal_iou


class SSNInstance:

    def __init__(self, start_frame, end_frame, video_frame_count,
                 fps=1, label=None,
                 best_iou=None, overlap_self=None):
        self.start_frame = start_frame
        self.end_frame = min(end_frame, video_frame_count)
        self._label = label
        self.fps = fps

        self.coverage = (end_frame - start_frame) / video_frame_count

        self.best_iou = best_iou
        self.overlap_self = overlap_self

        self.loc_reg = None
        self.size_reg = None

    def compute_regression_targets(self, gt_list, fg_thresh):
        if self.best_iou < fg_thresh:
            # background proposals do not need this
            return

        # find the groundtruth instance with the highest IOU
        ious = [temporal_iou((self.start_frame, self.end_frame), (gt.start_frame, gt.end_frame)) for gt in gt_list]
        best_gt_id = np.argmax(ious)

        best_gt = gt_list[best_gt_id]

        prop_center = (self.start_frame + self.end_frame) / 2
        gt_center = (best_gt.start_frame + best_gt.end_frame) / 2

        prop_size = self.end_frame - self.start_frame + 1
        gt_size = best_gt.end_frame - best_gt.start_frame + 1

        # get regression target:
        # (1). center shift propotional to the proposal duration
        # (2). logarithm of the groundtruth duration over proposal duraiton

        self.loc_reg = (gt_center - prop_center) / prop_size
        try:
            self.size_reg = math.log(gt_size / prop_size)
        except:
            print(gt_size, prop_size, self.start_frame, self.end_frame)
            raise

    @property
    def start_time(self):
        return self.start_frame / self.fps

    @property
    def end_time(self):
        return self.end_frame / self.fps

    @property
    def label(self):
        return self._label if self._label is not None else -1

    @property
    def regression_targets(self):
        return [self.loc_reg, self.size_reg] if self.loc_reg is not None else [0, 0]


class SSNVideoRecord:
    def __init__(self, prop_record):
        self._data = prop_record

        frame_count = int(self._data[1])

        # build instance record
        self.gt = [
            SSNInstance(int(x[1]), int(x[2]), frame_count, label=int(x[0]), best_iou=1.0) for x in self._data[2]
            if int(x[2]) > int(x[1])
        ]

        self.gt = list(filter(lambda x: x.start_frame < frame_count, self.gt))

        self.proposals = [
            SSNInstance(int(x[3]), int(x[4]), frame_count, label=int(x[0]),
                        best_iou=float(x[1]), overlap_self=float(x[2])) for x in self._data[3] if int(x[4]) > int(x[3])
        ]

        self.proposals = list(filter(lambda x: x.start_frame < frame_count, self.proposals))

    @property
    def id(self):
        return self._data[0]

    @property
    def num_frames(self):
        return int(self._data[1])

    def get_fg(self, fg_thresh, with_gt=True):
        fg = [p for p in self.proposals if p.best_iou > fg_thresh]
        if with_gt:
            fg.extend(self.gt)

        for x in fg:
            x.compute_regression_targets(self.gt, fg_thresh)
        return fg

    def get_negatives(self, incomplete_iou_thresh, bg_iou_thresh,
                      bg_coverage_thresh=0.01, incomplete_overlap_thresh=0.7):

        tag = [0] * len(self.proposals)

        incomplete_props = []
        background_props = []

        for i in range(len(tag)):
            if self.proposals[i].best_iou < incomplete_iou_thresh \
                    and self.proposals[i].overlap_self > incomplete_overlap_thresh:
                tag[i] = 1 # incomplete
                incomplete_props.append(self.proposals[i])

        for i in range(len(tag)):
            if tag[i] == 0 and \
                self.proposals[i].best_iou < bg_iou_thresh and \
                            self.proposals[i].coverage > bg_coverage_thresh:
                background_props.append(self.proposals[i])
        return incomplete_props, background_props


class SSNDataSet(data.Dataset):

    def __init__(self, root_path,
                 prop_file=None,
                 body_seg=5, aug_seg=2, video_centric=True,
                 new_length=1, modality='RGB',
                 image_tmpl='img_{:05d}.jpg', transform=None,
                 random_shift=True, test_mode=False,
                 prop_per_video=8, fg_ratio=1, bg_ratio=1, incomplete_ratio=6,
                 fg_iou_thresh=0.7,
                 bg_iou_thresh=0.01, incomplete_iou_thresh=0.3,
                 bg_coverage_thresh=0.02, incomplete_overlap_thresh=0.7,
                 gt_as_fg=True, reg_stats=None, test_interval=6, verbose=True,
                 exclude_empty=True, epoch_multiplier=1):

        self.root_path = root_path
        self.prop_file = prop_file
        self.verbose = verbose

        self.body_seg = body_seg
        self.aug_seg = aug_seg
        self.video_centric = video_centric
        self.exclude_empty = exclude_empty
        self.epoch_multiplier = epoch_multiplier

        self.new_length = new_length
        self.modality = modality
        self.image_tmpl = image_tmpl
        self.transform = transform
        self.random_shift = random_shift
        self.test_mode = test_mode
        self.test_interval = test_interval

        self.fg_iou_thresh = fg_iou_thresh
        self.incomplete_iou_thresh = incomplete_iou_thresh
        self.bg_iou_thresh = bg_iou_thresh

        self.bg_coverage_thresh = bg_coverage_thresh
        self.incomplete_overlap_thresh = incomplete_overlap_thresh

        self.starting_ratio = 0.5
        self.ending_ratio = 0.5

        self.gt_as_fg = gt_as_fg

        denum = fg_ratio + bg_ratio + incomplete_ratio

        self.fg_per_video = int(prop_per_video * (fg_ratio / denum))
        self.bg_per_video = int(prop_per_video * (bg_ratio / denum))
        self.incomplete_per_video = prop_per_video - self.fg_per_video - self.bg_per_video

        self._parse_prop_file(stats=reg_stats)

    def _load_image(self, directory, idx):
        if self.modality == 'RGB' or self.modality == 'RGBDiff':
            return [Image.open(os.path.join(directory, self.image_tmpl.format(idx))).convert('RGB')]
        elif self.modality == 'Flow':
            x_img = Image.open(os.path.join(directory, self.image_tmpl.format('x', idx))).convert('L')
            y_img = Image.open(os.path.join(directory, self.image_tmpl.format('y', idx))).convert('L')

            return [x_img, y_img]

    def _parse_prop_file(self, stats=None):
        prop_info = load_proposal_file(self.prop_file)

        self.video_list = [SSNVideoRecord(p) for p in prop_info]

        if self.exclude_empty:
            self.video_list = list(filter(lambda x: len(x.gt) > 0, self.video_list))

        self.video_dict = {v.id: v for v in self.video_list}

        # construct three pools:
        # 1. Foreground
        # 2. Background
        # 3. Incomplete

        self.fg_pool = []
        self.bg_pool = []
        self.incomp_pool = []

        for v in self.video_list:
            self.fg_pool.extend([(v.id, prop) for prop in v.get_fg(self.fg_iou_thresh, self.gt_as_fg)])

            incomp, bg = v.get_negatives(self.incomplete_iou_thresh, self.bg_iou_thresh,
                                         self.bg_coverage_thresh, self.incomplete_overlap_thresh)

            self.incomp_pool.extend([(v.id, prop) for prop in incomp])
            self.bg_pool.extend([(v.id, prop) for prop in bg])

        if stats is None:
            self._compute_regresssion_stats()
        else:
            self.stats = stats

        if self.verbose:
            print("""
            
            SSNDataset: Proposal file {prop_file} parsed.
            
            There are {pnum} usable proposals from {vnum} videos.
            {fnum} foreground proposals
            {inum} incomplete_proposals
            {bnum} background_proposals
            
            Sampling config:
            FG/BG/INC: {fr}/{br}/{ir}
            Video Centric: {vc}
            
            Epoch size multiplier: {em}
            
            Regression Stats:
            Location: mean {stats[0][0]:.05f} std {stats[1][0]:.05f}
            Duration: mean {stats[0][1]:.05f} std {stats[1][1]:.05f}
            """.format(prop_file=self.prop_file, pnum=len(self.fg_pool) + len(self.bg_pool) + len(self.incomp_pool),
                       fnum=len(self.fg_pool), inum=len(self.incomp_pool), bnum=len(self.bg_pool),
                       fr=self.fg_per_video, br=self.bg_per_video, ir=self.incomplete_per_video, vnum=len(self.video_dict),
                       vc=self.video_centric, stats=self.stats, em=self.epoch_multiplier))
        else:
            print("""
                        SSNDataset: Proposal file {prop_file} parsed.   
            """.format(prop_file=self.prop_file))


    def _video_centric_sampling(self, video):

        fg = video.get_fg(self.fg_iou_thresh, self.gt_as_fg)
        incomp, bg = video.get_negatives(self.incomplete_iou_thresh, self.bg_iou_thresh,
                                     self.bg_coverage_thresh, self.incomplete_overlap_thresh)

        def sample_video_proposals(proposal_type, video_id, video_pool, requested_num, dataset_pool):
            if len(video_pool) == 0:
                # if there is nothing in the video pool, go fetch from the dataset pool
                return [(dataset_pool[x], proposal_type) for x in np.random.choice(len(dataset_pool), requested_num, replace=False)]
            else:
                replicate = len(video_pool) < requested_num
                idx = np.random.choice(len(video_pool), requested_num, replace=replicate)
                return [((video_id, video_pool[x]), proposal_type) for x in idx]

        out_props = []
        out_props.extend(sample_video_proposals(0, video.id, fg, self.fg_per_video, self.fg_pool))  # sample foreground
        out_props.extend(sample_video_proposals(1, video.id, incomp, self.incomplete_per_video, self.incomp_pool))  # sample incomp.
        out_props.extend(sample_video_proposals(2, video.id, bg, self.bg_per_video, self.bg_pool))  # sample background

        return out_props

    def _random_sampling(self):
        out_props = []

        out_props.extend([(x, 0) for x in np.random.choice(self.fg_pool, self.fg_per_video, replace=False)])
        out_props.extend([(x, 1) for x in np.random.choice(self.incomp_pool, self.incomplete_per_video, replace=False)])
        out_props.extend([(x, 2) for x in np.random.choice(self.bg_pool, self.bg_per_video, replace=False)])

        return out_props

    def _sample_indices(self, valid_length, num_seg):
        """

        :param record: VideoRecord
        :return: list
        """

        average_duration = (valid_length + 1) // num_seg
        if average_duration > 0:
            # normal cases
            offsets = np.multiply(list(range(num_seg)), average_duration) \
                             + randint(average_duration, size=num_seg)
        elif valid_length > num_seg:
            offsets = np.sort(randint(valid_length, size=num_seg))
        else:
            offsets = np.zeros((num_seg, ))

        return offsets

    def _get_val_indices(self, valid_length, num_seg):

        if valid_length > num_seg:
            tick = valid_length / float(num_seg)
            offsets = np.array([int(tick / 2.0 + tick * x) for x in range(num_seg)])
        else:
            offsets = np.zeros((num_seg,))

        return offsets

    def _sample_ssn_indices(self, prop, frame_cnt):
        start_frame = prop.start_frame + 1
        end_frame = prop.end_frame

        duration = end_frame - start_frame + 1
        assert duration != 0, (prop.start_frame, prop.end_frame, prop.best_iou)
        valid_length = duration - self.new_length

        valid_starting = max(1, start_frame - int(duration * self.starting_ratio))
        valid_ending = min(frame_cnt - self.new_length + 1, end_frame + int(duration * self.ending_ratio))

        valid_starting_length = (start_frame - valid_starting - self.new_length + 1)
        valid_ending_length = (valid_ending - end_frame - self.new_length + 1)

        starting_scale = (valid_starting_length + self.new_length - 1) / (duration * self.starting_ratio)
        ending_scale = (valid_ending_length + self.new_length - 1) / (duration * self.ending_ratio)

        # get starting
        starting_offsets = (self._sample_indices(valid_starting_length, self.aug_seg) if self.random_shift
                            else self._get_val_indices(valid_starting_length, self.aug_seg)) + valid_starting
        course_offsets = (self._sample_indices(valid_length, self.body_seg) if self.random_shift
                          else self._get_val_indices(valid_length, self.body_seg)) + start_frame
        ending_offsets = (self._sample_indices(valid_ending_length, self.aug_seg) if self.random_shift
                          else self._get_val_indices(valid_ending_length, self.aug_seg)) + end_frame

        offsets = np.concatenate((starting_offsets, course_offsets, ending_offsets))
        stage_split = [self.aug_seg, self.aug_seg + self.body_seg, self.aug_seg * 2 + self.body_seg]
        return offsets, starting_scale, ending_scale, stage_split

    def _load_prop_data(self, prop):

        # read frame count
        frame_cnt = self.video_dict[prop[0][0]].num_frames

        # sample segment indices
        prop_indices, starting_scale, ending_scale, stage_split = self._sample_ssn_indices(prop[0][1], frame_cnt)

        # turn prop into standard format

        # get label
        if prop[1] == 0:
            label = prop[0][1].label
        elif prop[1] == 1:
            label = prop[0][1].label  # incomplete
        elif prop[1] == 2:
            label = 0  # background
        else:
            raise ValueError()
        frames = []
        for idx, seg_ind in enumerate(prop_indices):
            p = int(seg_ind)
            for x in range(self.new_length):
                frames.extend(self._load_image(prop[0][0], min(frame_cnt, p+x)))

        # get regression target
        if prop[1] == 0:
            reg_targets = prop[0][1].regression_targets
            reg_targets = (reg_targets[0] - self.stats[0][0]) / self.stats[1][0], \
                          (reg_targets[1] - self.stats[0][1]) / self.stats[1][1]
        else:
            reg_targets = (0.0, 0.0)

        return frames, label, reg_targets, starting_scale, ending_scale, stage_split, prop[1]

    def _compute_regresssion_stats(self):
        if self.verbose:
            print("computing regression target normalizing constants")
        targets = []
        for video in self.video_list:
            fg = video.get_fg(self.fg_iou_thresh, False)
            for p in fg:
                targets.append(list(p.regression_targets))

        self.stats = np.array((np.mean(targets, axis=0), np.std(targets, axis=0)))

    def get_test_data(self, video, test_interval, gen_batchsize=4):
        props = video.proposals
        video_id = video.id
        frame_cnt = video.num_frames
        frame_ticks = np.arange(0, frame_cnt - self.new_length, test_interval, dtype=np.int) + 1

        num_sampled_frames = len(frame_ticks)

        # avoid empty proposal list
        if len(props) == 0:
            props.append(SSNInstance(0, frame_cnt - 1, frame_cnt))

        # process proposals to subsampled sequences
        rel_prop_list = []
        proposal_tick_list = []
        scaling_list = []
        for proposal in props:
            rel_prop = proposal.start_frame / frame_cnt, proposal.end_frame / frame_cnt
            rel_duration = rel_prop[1] - rel_prop[0]
            rel_starting_duration = rel_duration * self.starting_ratio
            rel_ending_duration = rel_duration * self.ending_ratio
            rel_starting = rel_prop[0] - rel_starting_duration
            rel_ending = rel_prop[1] + rel_ending_duration

            real_rel_starting = max(0.0, rel_starting)
            real_rel_ending = min(1.0, rel_ending)

            starting_scaling = (rel_prop[0] - real_rel_starting) / rel_starting_duration
            ending_scaling = (real_rel_ending - rel_prop[1]) / rel_ending_duration

            proposal_ticks = int(real_rel_starting * num_sampled_frames), int(rel_prop[0] * num_sampled_frames), \
                             int(rel_prop[1] * num_sampled_frames), int(real_rel_ending * num_sampled_frames)

            rel_prop_list.append(rel_prop)
            proposal_tick_list.append(proposal_ticks)
            scaling_list.append((starting_scaling, ending_scaling))

        # load frames
        # Since there are many frames for each video during testing, instead of returning the read frames,
        # we return a generator which gives the frames in small batches, this lower the memory burden
        # and runtime overhead. Usually setting batchsize=4 would fit most cases.
        def frame_gen(batchsize):
            frames = []
            cnt = 0
            for idx, seg_ind in enumerate(frame_ticks):
                p = int(seg_ind)
                for x in range(self.new_length):
                    frames.extend(self._load_image(video_id, min(frame_cnt, p+x)))
                cnt += 1

                if cnt % batchsize == 0:
                    frames = self.transform(frames)
                    yield frames
                    frames = []

            if len(frames):
                frames = self.transform(frames)
                yield frames

        return frame_gen(gen_batchsize), len(frame_ticks), torch.from_numpy(np.array(rel_prop_list)), \
               torch.from_numpy(np.array(proposal_tick_list)), torch.from_numpy(np.array(scaling_list))

    def get_training_data(self, index):
        if self.video_centric:
            video = self.video_list[index]
            props = self._video_centric_sampling(video)
        else:
            props = self._random_sampling()

        out_frames = []
        out_prop_len = []
        out_prop_scaling = []
        out_prop_type = []
        out_prop_labels = []
        out_prop_reg_targets = []
        out_stage_split = []
        for idx, p in enumerate(props):
            prop_frames, prop_label, reg_targets, starting_scale, ending_scale, stage_split, prop_type = self._load_prop_data(
                p)

            processed_frames = self.transform(prop_frames)
            out_frames.append(processed_frames)
            out_prop_len.append(self.body_seg + 2 * self.aug_seg)
            out_prop_scaling.append([starting_scale, ending_scale])
            out_prop_labels.append(prop_label)
            out_prop_reg_targets.append(reg_targets)
            out_prop_type.append(prop_type)
            out_stage_split.append(stage_split)

        out_prop_len = torch.from_numpy(np.array(out_prop_len))
        out_prop_scaling = torch.from_numpy(np.array(out_prop_scaling, dtype=np.float32))
        out_prop_labels = torch.from_numpy(np.array(out_prop_labels))
        out_prop_reg_targets = torch.from_numpy(np.array(out_prop_reg_targets, dtype=np.float32))
        out_prop_type = torch.from_numpy(np.array(out_prop_type))
        out_stage_split = torch.from_numpy(np.array(out_stage_split))
        out_frames = torch.cat(out_frames)
        return out_frames, out_prop_len, out_prop_scaling, out_prop_type, out_prop_labels, \
               out_prop_reg_targets, out_stage_split

    def get_all_gt(self):
        gt_list = []
        for video in self.video_list:
            vid = video.id
            gt_list.extend([[vid, x.label - 1, x.start_frame / video.num_frames,
                             x.end_frame / video.num_frames] for x in video.gt])
        return gt_list

    def __getitem__(self, index):
        real_index = index % len(self.video_list)
        if self.test_mode:
            return self.get_test_data(self.video_list[real_index], self.test_interval)
        else:
            return self.get_training_data(real_index)

    def __len__(self):
        return len(self.video_list) * self.epoch_multiplier