187 lines
7.4 KiB
Python
187 lines
7.4 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import copy
|
|
import random
|
|
|
|
import decord
|
|
import numpy as np
|
|
import torch
|
|
from detectron2.data.transforms import (ExtentTransform, RandomBrightness,
|
|
RandomFlip, ResizeShortestEdge)
|
|
from detectron2.structures import Boxes, Instances
|
|
from scipy.interpolate import interp1d
|
|
|
|
|
|
def inp_boxes(boxes: dict, start, end):
|
|
idxs = sorted([int(i) for i in boxes.keys()])
|
|
bbox = [boxes[str(i)] for i in idxs]
|
|
new_bboxes = []
|
|
for i in range(4):
|
|
f = interp1d(idxs, [b[i] for b in bbox])
|
|
new_b = f(list(range(start, end + 1)))
|
|
new_bboxes.append(new_b)
|
|
new_bboxes = np.stack(new_bboxes, axis=1)
|
|
return new_bboxes
|
|
|
|
|
|
def assign_label(start, end, data_dict):
|
|
"""
|
|
根据视频起始位置，以及标注的label,给这小段视频安排bbox检测标签
|
|
方法，取交集，交集占到样本的一半或者标签的一半，即将该label赋给样本
|
|
:param start: 起始帧号（含）
|
|
:param end: 结束帧号（含）
|
|
:param labels: 标注的label, 字符串形式
|
|
:return:[[行为，x1,y1,x2,y2],]
|
|
"""
|
|
if 'actions' not in data_dict:
|
|
return []
|
|
scale = data_dict['scale']
|
|
gt_labels = []
|
|
for action in data_dict['actions']:
|
|
low = max(int(action['start']), start)
|
|
high = min(int(action['end']), end)
|
|
inter = 0 if low > high else high - low
|
|
if inter > (end - start) * 0.7 or inter > (action['end']
|
|
- action['start']) * 0.7:
|
|
boxes = inp_boxes(action['boxes'], low, high)
|
|
box = boxes.mean(axis=0) / scale
|
|
label = [action['label']] + box.tolist()
|
|
gt_labels.append(label)
|
|
return gt_labels
|
|
|
|
|
|
class VideoDetMapper:
|
|
|
|
def __init__(self,
|
|
classes_id_map,
|
|
used_seconds=2,
|
|
input_frames=4,
|
|
is_train=True,
|
|
tile=False):
|
|
self.classes_id = classes_id_map
|
|
self.is_train = is_train
|
|
self.used_seconds = used_seconds
|
|
self.input_frames = input_frames
|
|
self.tile = tile
|
|
self.trans = [RandomBrightness(0.5, 1.5)]
|
|
self.tfm_gens = [
|
|
ResizeShortestEdge((480, 512, 544, 576, 608, 640, 672, 704, 736,
|
|
768) if is_train else 512,
|
|
1280 if is_train else 896, 'choice')
|
|
]
|
|
if is_train:
|
|
self.tfm_gens.append(RandomFlip())
|
|
|
|
def __call__(self, data_dict):
|
|
data_dict = copy.deepcopy(data_dict)
|
|
try:
|
|
data_dict = self._call(data_dict)
|
|
except Exception as e:
|
|
print(data_dict['path:FILE'], e)
|
|
data_dict = None
|
|
return data_dict
|
|
|
|
def _call(self, data_dict):
|
|
video_name = data_dict['path:FILE']
|
|
if data_dict['actions'] is not None:
|
|
data_dict['actions'] = eval(data_dict['actions'])
|
|
else:
|
|
data_dict['actions'] = []
|
|
|
|
v = decord.VideoReader(video_name, ctx=decord.cpu(0))
|
|
num_frames = len(v)
|
|
used_frames = max(int((1 + random.random()) * v.get_avg_fps()), 1)
|
|
if self.is_train:
|
|
start_idx = random.randint(0, max(0, num_frames - used_frames))
|
|
else:
|
|
start_idx = max(0, num_frames - used_frames) // 2
|
|
idxs = np.linspace(start_idx, min(start_idx + used_frames, num_frames) - 1, self.input_frames) \
|
|
.round().astype('int32').tolist()
|
|
imgs = v.get_batch(idxs).asnumpy()
|
|
del v
|
|
labels = assign_label(idxs[0], idxs[-1] + 1, data_dict)
|
|
bboxes = np.array([label[-4:] for label in labels])
|
|
|
|
if self.is_train:
|
|
if self.tile:
|
|
imgs, labels, bboxes = self.random_tile(
|
|
video_name, imgs, labels, bboxes, pos_choices=[1, 1, 2, 4])
|
|
else:
|
|
imgs, labels, bboxes = self.random_tile(
|
|
video_name, imgs, labels, bboxes, pos_choices=[1])
|
|
|
|
for g in self.trans:
|
|
tfm = g.get_transform(imgs)
|
|
imgs = tfm.apply_image(imgs)
|
|
imgs, bboxes = self.random_extent(imgs, bboxes)
|
|
|
|
for trans in self.tfm_gens:
|
|
tfm = trans.get_transform(imgs[0])
|
|
imgs = np.stack([tfm.apply_image(img) for img in imgs])
|
|
bboxes = tfm.apply_box(bboxes)
|
|
|
|
_, h, w, c = imgs.shape
|
|
data_dict['height'] = h
|
|
data_dict['width'] = w
|
|
gt_boxes = Boxes(torch.from_numpy(bboxes)) # XYXY_ABS
|
|
gt_classes = [self.classes_id[label[0]]
|
|
for label in labels] # N is background
|
|
instances = Instances((data_dict['height'], data_dict['width']))
|
|
instances.set('gt_boxes', gt_boxes)
|
|
instances.set('gt_classes',
|
|
torch.as_tensor(gt_classes, dtype=torch.int64))
|
|
data_dict['instances'] = instances
|
|
data_dict['frames'] = torch.as_tensor(
|
|
np.ascontiguousarray(imgs.transpose([3, 0, 1, 2])))
|
|
return data_dict
|
|
|
|
def random_tile(self, name, imgs, labels, bboxes,
|
|
pos_choices=(1, 1, 2, 4)):
|
|
_, h, w, c = imgs.shape
|
|
bboxes = bboxes.tolist()
|
|
if len(labels) == 0: # 负样本 1/2, 1, 2, 4
|
|
ratio = random.choice([0, 1, 2, 4])
|
|
if ratio == 0: # 随机取部分区域
|
|
h0, w0 = random.randint(0, h // 2), random.randint(0, w // 2)
|
|
imgs = imgs[:, h0:h0 + h // 2, w0:w0 + h // 2]
|
|
elif ratio == 2:
|
|
imgs = np.tile(imgs,
|
|
(1, 1, 2,
|
|
1)) if h > w else np.tile(imgs, (1, 2, 1, 1))
|
|
elif ratio == 4:
|
|
imgs = np.tile(imgs, (1, 2, 2, 1))
|
|
else: # 正样本 1, 2, 4
|
|
ratio = random.choice(pos_choices)
|
|
if ratio == 2:
|
|
labels = labels * 2
|
|
if h >= w: # 左右拼接
|
|
imgs = np.tile(imgs, (1, 1, 2, 1))
|
|
bbox2 = [[x1 + w, y1, x2 + w, y2]
|
|
for x1, y1, x2, y2 in bboxes]
|
|
else: # 上下拼接
|
|
imgs = np.tile(imgs, (1, 2, 1, 1))
|
|
bbox2 = [[x1, y1 + h, x2, y2 + h]
|
|
for x1, y1, x2, y2 in bboxes]
|
|
bboxes = bboxes + bbox2
|
|
elif ratio == 4:
|
|
labels = labels * 4
|
|
imgs = np.tile(imgs, (1, 2, 2, 1))
|
|
bbox2 = [[x1 + w, y1, x2 + w, y2] for x1, y1, x2, y2 in bboxes] + \
|
|
[[x1, y1 + h, x2, y2 + h] for x1, y1, x2, y2 in bboxes] + \
|
|
[[x1 + w, y1 + h, x2 + w, y2 + h] for x1, y1, x2, y2 in bboxes]
|
|
bboxes = bboxes + bbox2
|
|
bboxes = np.array(bboxes)
|
|
return imgs.copy(), labels, bboxes
|
|
|
|
def random_extent(self, imgs, bboxes):
|
|
t, h, w, c = imgs.shape
|
|
r_h, r_w = int(h * 0.1), int(w * 0.1)
|
|
x0, y0 = random.randint(-r_w, r_w), random.randint(-r_h, r_h)
|
|
x1, y1 = random.randint(w - r_w,
|
|
w + r_w), random.randint(h - r_h, h + r_h)
|
|
tfm = ExtentTransform((x0, y0, x1, y1), output_size=(y1 - y0, x1 - x0))
|
|
imgs = np.stack([tfm.apply_image(img) for img in imgs])
|
|
bboxes = tfm.apply_box(bboxes)
|
|
return imgs, bboxes
|