Spaces:
Running
on
T4
Running
on
T4
File size: 3,196 Bytes
d4b77ac |
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 |
# --------------------------------------------------------
# SiamMask
# Licensed under The MIT License
# Written by Qiang Wang (wangqiang2015 at ia.ac.cn)
# --------------------------------------------------------
from os.path import join
from os import listdir
import json
import numpy as np
print('load json (raw vid info), please wait 20 seconds~')
vid = json.load(open('vid.json', 'r'))
def check_size(frame_sz, bbox):
min_ratio = 0.1
max_ratio = 0.75
# only accept objects >10% and <75% of the total frame
area_ratio = np.sqrt((bbox[2]-bbox[0])*(bbox[3]-bbox[1])/float(np.prod(frame_sz)))
ok = (area_ratio > min_ratio) and (area_ratio < max_ratio)
return ok
def check_borders(frame_sz, bbox):
dist_from_border = 0.05 * (bbox[2] - bbox[0] + bbox[3] - bbox[1])/2
ok = (bbox[0] > dist_from_border) and (bbox[1] > dist_from_border) and \
((frame_sz[0] - bbox[2]) > dist_from_border) and \
((frame_sz[1] - bbox[3]) > dist_from_border)
return ok
snippets = dict()
n_snippets = 0
n_videos = 0
for subset in vid:
for video in subset:
n_videos += 1
frames = video['frame']
id_set = []
id_frames = [[]] * 60 # at most 60 objects
for f, frame in enumerate(frames):
objs = frame['objs']
frame_sz = frame['frame_sz']
for obj in objs:
trackid = obj['trackid']
occluded = obj['occ']
bbox = obj['bbox']
# if occluded:
# continue
#
# if not(check_size(frame_sz, bbox) and check_borders(frame_sz, bbox)):
# continue
#
# if obj['c'] in ['n01674464', 'n01726692', 'n04468005', 'n02062744']:
# continue
if trackid not in id_set:
id_set.append(trackid)
id_frames[trackid] = []
id_frames[trackid].append(f)
if len(id_set) > 0:
snippets[video['base_path']] = dict()
for selected in id_set:
frame_ids = sorted(id_frames[selected])
sequences = np.split(frame_ids, np.array(np.where(np.diff(frame_ids) > 1)[0]) + 1)
sequences = [s for s in sequences if len(s) > 1] # remove isolated frame.
for seq in sequences:
snippet = dict()
for frame_id in seq:
frame = frames[frame_id]
for obj in frame['objs']:
if obj['trackid'] == selected:
o = obj
continue
snippet[frame['img_path'].split('.')[0]] = o['bbox']
snippets[video['base_path']]['{:02d}'.format(selected)] = snippet
n_snippets += 1
print('video: {:d} snippets_num: {:d}'.format(n_videos, n_snippets))
train = {k:v for (k,v) in snippets.items() if 'train' in k}
val = {k:v for (k,v) in snippets.items() if 'val' in k}
json.dump(train, open('train.json', 'w'), indent=4, sort_keys=True)
json.dump(val, open('val.json', 'w'), indent=4, sort_keys=True)
print('done!')
|