def main():
"""
Main function to spawn the train and test process.
"""
args = parse_args()
cfg = load_config(args)
#cfg.DEMO.WEBCAM = 0
cfg.DEMO.WEBCAM = -1
cfg.DEMO.INPUT_VIDEO = "demo_test/demo_in2.mp4"
cfg.NUM_GPUS = 1
cfg.TRAIN.ENABLE = False
cfg.TEST.ENABLE = False
cfg.DEMO.OUTPUT_FILE = "demo_test/demo_out2.mp4"
cfg.DEMO.ENABLE = True
# Perform training.
if cfg.TRAIN.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=train)
# Perform multi-clip testing.
if cfg.TEST.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=test)
# Perform model visualization.
if cfg.TENSORBOARD.ENABLE and cfg.TENSORBOARD.MODEL_VIS.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=visualize)
# Run demo.
if cfg.DEMO.ENABLE:
demo(cfg)
def main():
args = parse_args()
cfg = load_config(args)
launch_job(
cfg=cfg, init_method=args.init_method, func=benchmark_data_loading
)
def main():
"""
Main function to spawn the train and test process.
"""
args = parse_args()
cfg = load_config(args)
# Perform training.
if cfg.TRAIN.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=train)
# Perform multi-clip testing.
if cfg.TEST.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=test)
# Perform model visualization.
if cfg.TENSORBOARD.ENABLE and (
cfg.TENSORBOARD.MODEL_VIS.ENABLE
or cfg.TENSORBOARD.WRONG_PRED_VIS.ENABLE
):
launch_job(cfg=cfg, init_method=args.init_method, func=visualize)
# Run demo.
if cfg.DEMO.ENABLE:
demo(cfg)
def export_pytorch_model():
weights = 'checkpoints/checkpoint_epoch_00028.pyth'
args = parse_args()
cfg = load_config(args)
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
model = build_model(cfg).to(device)
# chkpt = torch.load(weights, map_location="cpu")
chkpt = torch.load(weights, map_location=device)
model.load_state_dict(chkpt['model_state'])
# try:
# model_dict = model.module.state_dict()
# except AttributeError:
# # 读取原始状态及参数， ## 多GPU训练，导致训练存储的模型时key会加上model
# model_dict = model.state_dict()
# # 将pretrained_dict里不属于model_dict的键剔除掉
# chkpt = {k: v for k, v in chkpt.items() if k in model_dict}
# print("load pretrain model")
# model_dict.update(chkpt)
# # model.state_dict(model_dict)
# model.load_state_dict(model_dict)
# z转换为评估模型
model.eval()
# e1 = torch.rand(1, 3, 8, 224, 224).cuda()
# e2 = torch.rand(1, 3, 32, 224, 224).cuda()
e1 = torch.rand(8, 3, 256, 455) # .fill_(0)
e2 = torch.rand(32, 3, 256, 455) # .fill_(0)
e3 = [e1, e2]
# e4 = torch.rand(1, 5).cuda()
# cuda()
e4 = torch.rand(1, 1, 1, 5) # .fill_(0) rand(1, 1, 3, 5)
import numpy
import numpy as np
numpy.save("input00.npy", e3[0].numpy())
numpy.save("input11.npy", e3[1].numpy())
numpy.save("input22.npy", e4.numpy())
input0 = torch.from_numpy(np.load("input0.npy"))
input1 = torch.from_numpy(np.load("input1.npy"))
input2 = torch.from_numpy(np.load("input2.npy"))
pred = model(e3, e4)
print(pred)
# exit(0)
input3 = [input0, input1]
# traced_script_module = torch.jit.trace(model, (e3, e4))
traced_script_module = torch.jit.trace(model, (input3, input2))
# print(traced_script_module.graph)
print(traced_script_module(input3, input2)) # .forward
traced_script_module.save("weights/sf18_pytorch_cpu4503.pt")
print("out put save")
def main():
"""
Main function to spawn the train and test process.
"""
args = parse_args()
cfg = load_config(args)
if cfg.DEMO.ENABLE:
original_num_gpus = cfg.NUM_GPUS
# Set num_gpus to 1 for the demo
cfg.NUM_GPUS = 1
launch_job(cfg=cfg, init_method=args.init_method, func=run_demo)
# Set num gpus back to original
cfg.NUM_GPUS = original_num_gpus
# Perform training.
if cfg.TRAIN.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=train)
# Perform multi-clip testing.
if cfg.TEST.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=test)
if cfg.DEMO_ORIGINAL.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=demo_original)
if cfg.TENSORBOARD.ENABLE and cfg.TENSORBOARD.MODEL_VIS.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=visualize)
def main():
"""
Main function to spawn the train and test process.
"""
args = parse_args()
cfg = load_config(args)
if cfg.TEST.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=test)
def main():
"""
Main function to spawn the train and test process.
"""
args = parse_args()
cfg = load_config(args)
# Perform training.
if cfg.TRAIN.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=train)
# Perform multi-clip testing.
if cfg.TEST.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=test)
def __init__(
self,
num_classes,
class_names_path,
top_k=1,
colormap="rainbow",
thres=0.7,
lower_thres=0.3,
common_class_names=None,
mode="top-k",
):
"""
Args:
num_classes (int): total number of classes.
class_names_path (str): path to json file that maps class names to ids.
Must be in the format {classname: id}.
top_k (int): number of top predicted classes to plot.
colormap (str): the colormap to choose color for class labels from.
See https://matplotlib.org/tutorials/colors/colormaps.html
thres (float): threshold for picking predicted classes to visualize.
lower_thres (Optional[float]): If `common_class_names` if given,
this `lower_thres` will be applied to uncommon classes and
`thres` will be applied to classes in `common_class_names`.
common_class_names (Optional[list of str(s)]): list of common class names
to apply `thres`. Class names not included in `common_class_names` will
have `lower_thres` as a threshold. If None, all classes will have `thres` as a threshold.
This is helpful for model trained on highly imbalanced dataset.
mode (str): Supported modes are {"top-k", "thres"}.
This is used for choosing predictions for visualization.
"""
assert mode in ["top-k",
"thres"], "Mode {} is not supported.".format(mode)
self.mode = mode
self.num_classes = num_classes
self.class_names, _, _ = get_class_names(class_names_path, None, None)
self.top_k = top_k
self.thres = thres
self.lower_thres = lower_thres
if mode == "thres":
self._get_thres_array(common_class_names=common_class_names)
self.color_map = plt.get_cmap(colormap)
args = parse_args()
cfg = load_config(args)
self.tracker = HungarianTracker(alpha=cfg.DEMO.ALPHA)
def main():
"""
Main function to spawn the train and test process.
"""
args = parse_args()
cfg = load_config(args)
# Perform training.
if cfg.TRAIN.ENABLE:
if cfg.NUM_GPUS > 1:
torch.multiprocessing.spawn(
mpu.run,
nprocs=cfg.NUM_GPUS,
args=(
cfg.NUM_GPUS,
train,
args.init_method,
cfg.SHARD_ID,
cfg.NUM_SHARDS,
cfg.DIST_BACKEND,
cfg,
),
daemon=False,
)
else:
train(cfg=cfg)
# Perform multi-clip testing.
if cfg.TEST.ENABLE:
if cfg.NUM_GPUS > 1:
torch.multiprocessing.spawn(
mpu.run,
nprocs=cfg.NUM_GPUS,
args=(
cfg.NUM_GPUS,
test,
args.init_method,
cfg.SHARD_ID,
cfg.NUM_SHARDS,
cfg.DIST_BACKEND,
cfg,
),
daemon=False,
)
else:
test(cfg=cfg)
def main():
args = parse_args()
cfg = load_config(args)
# Perform training.
if cfg.TRAIN.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=train)
# Perform multi-clip testing.
if cfg.TEST.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=test)
# Perform feature extraction.
if cfg.MODEL.EXTRACTOR:
extract(cfg)
if cfg.MODEL.VIDEO_EXTRACTOR:
video_extract(cfg)
def main():
"""
Main function to spawn the train and test process.
"""
args = parse_args()
cfg = load_config(args)
if cfg.TRAIN.ENABLE:
if cfg.TRAIN.ONLY_DES:
launch_job(cfg=cfg, init_method=args.init_method, func=train_des)
else:
launch_job(cfg=cfg, init_method=args.init_method, func=train)
else:
if cfg.TRAIN.ONLY_DES:
launch_job(cfg=cfg,
init_method=args.init_method,
func=test_implementation_des)
else:
launch_job(cfg=cfg,
init_method=args.init_method,
func=test_implementation)
def main():
args = parse_args()
cfg = load_config(args)
if cfg.NUM_GPUS > 1:
torch.multiprocessing.spawn(
mpu.run,
nprocs=cfg.NUM_GPUS,
args=(
cfg.NUM_GPUS,
benchmark_data,
args.init_method,
cfg.SHARD_ID,
cfg.NUM_SHARDS,
cfg.DIST_BACKEND,
cfg,
),
daemon=False,
)
else:
benchmark_data(cfg=cfg)
def main():
"""
Main function to spawn the train and test process.
"""
args = parse_args()
cfg = load_config(args)
uuid = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, uuid)
os.makedirs(cfg.OUTPUT_DIR)
# Perform training.
if cfg.TRAIN.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=train)
# # Perform multi-clip testing.
# if cfg.TEST.ENABLE:
# launch_job(cfg=cfg, init_method=args.init_method, func=test)
# if cfg.DEMO.ENABLE:
# launch_job(cfg=cfg, init_method=args.init_method, func=demo)
if cfg.TENSORBOARD.ENABLE and cfg.TENSORBOARD.MODEL_VIS.ENABLE:
launch_job(cfg=cfg, init_method=args.init_method, func=visualize)
# Save gif
if (i+1) % 5 == 0:
layer_path = self.dir_path + 'layer/'
if not os.path.exists(layer_path):
os.mkdir(layer_path)
path = layer_path + str(self.model_name) + '_iter' + str(i) + '_path' + str(j) + '_' + \
str(self.layer) + '_f' + str(self.filter) + '_lr' + str(self.initial_learning_rate) + "_wd" \
+ str(self.weight_decay)
save_gif(created_video, path, stream_type="rgb")
if __name__ == '__main__':
args = parse_args()
cfg = load_config(args)
# Construct the model
model = build_model(cfg)
load_checkpoint(checkpoint_path, model, data_parallel=False, optimizer=None, inflation=False, convert_from_caffe2=True)
cnn_layer = "s2.pathway0_res0.branch1" # "conv3d_0c_1x1.conv3d"
filter_pos = 0
device = torch.device('cuda:0')
model = model.to(device)
layer_vis = CNNLayerVisualization(model, cnn_layer, filter_pos, device)
# Layer visualization with pytorch hooks
layer_vis.visualise_layer_with_hooks()
bbox['bottom_right_x'], bbox['bottom_right_y'])
reduced_tasks = {}
for pred in tasks:
bbox = transform_bbox(pred['box'])
if pred['task_id'] not in reduced_tasks:
reduced_tasks[pred['task_id']] = [bbox]
else:
reduced_tasks[pred['task_id']].append(bbox)
return reduced_tasks
def main(predictions_path):
"""
Performs a hungarian traking on each tasks
"""
with open(predictions_path, 'r') as json_file:
data = json.load(json_file)
data = reduce_tasks(data)
tracker = HungarianTracker()
for task_id in tqdm(data):
ids = tracker.advance(data[task_id])
print(f">> task {task_id} ids = {ids}")
if __name__ == "__main__":
ARGS = parse_args()
CFG = load_config(ARGS)
PREDICTIONS_PATH = CFG.DEMO.PREDICTIONS_FILE_PATH
main(PREDICTIONS_PATH)
def get_model(
num_classes,
model_name='TSN',
freeze_layers=True,
model_path=None,
):
class_counts = (125, 352)
segment_count = 8
base_model = 'resnet50'
if model_name == 'BLV':
backbone_setting = {
'dropout': 0.5,
'pretrained': True,
'alpha': 2,
'depth': 101,
'beta': 4,
'input_channels': 3,
'num_classes': 174,
'dataset': 'st2stv2',
'groups': 64,
'imagenet_blnet_pretrained': False,
'blending_frames': 3
}
net = bLVNet_TAM(backbone_setting)
elif model_name == 'slowfast':
args = parse_args()
cfg = load_config(args)
cfg.NUM_GPUS = 1
cfg.TRAIN.CHECKPOINT_FILE_PATH = "SLOWFAST_4x16_R50.pkl"
net = build_model(cfg)
elif model_name == 'TSM':
net = get_tsm(num_classes)
# net = torch.hub.load(repo, 'TSM', class_counts, segment_count, 'RGB',
# base_model=base_model,
# pretrained='epic-kitchens')
elif model_name == 'I3D':
net = resnet.i3_res50(400)
else:
net = torch.hub.load(repo,
model_name,
class_counts,
segment_count,
'RGB',
base_model=base_model,
pretrained='epic-kitchens',
force_reload=True)
if freeze_layers:
for param in net.parameters():
param.requires_grad = False
if model_name == 'TSN': # or model_name == 'TSM':
net.fc_verb = torch.nn.Linear(2048, num_classes)
elif model_name == 'TRN':
net.consensus.classifiers[0] = torch.nn.Linear(512, num_classes)
elif model_name == 'BLV':
net.new_fc = torch.nn.Linear(2048, num_classes)
elif model_name == 'slowfast':
net.head.projection = torch.nn.Linear(2304, num_classes)
net.head.act = None
elif model_name == 'I3D':
net.fc = torch.nn.Linear(2048, num_classes)
elif model_name == 'TSM':
net.new_fc = torch.nn.Linear(2048, num_classes)
if model_path is not None:
load_model(net, model_path)
if model_name == 'BLV':
for param in net.baseline_model.layer4.parameters():
param.requires_grad = True
elif model_name in ['TRN', 'TSN', 'TSM']:
for param in net.base_model.layer4.parameters():
param.requires_grad = True
elif model_name == 'I3D':
for param in net.layer4.parameters():
param.requires_grad = True
use_gpu = torch.cuda.is_available()
if use_gpu:
net = net.cuda()
return net
def draw_clip(
self,
frames,
preds,
bboxes=None,
nbboxes=None,
text_alpha=0.5,
ground_truth=False,
keyframe_idx=None,
repeat_frame=1,
):
"""
Draw predicted labels or ground truth classes to clip. Draw bouding boxes to clip
if bboxes is provided. Boxes will gradually fade in and out the clip, centered around
the clip's central frame.
Args:
frames (array-like): video data in the shape (T, H, W, C).
preds (tensor): a tensor of shape (num_boxes, num_classes) that contains all of the confidence scores
of the model. For recognition task or for ground_truth labels, input shape can be (num_classes,).
bboxes (Optional[tensor]): shape (num_boxes, 4) that contains the coordinates of the bounding boxes.
nbboxes (Optional[tensor]): shape (num_boxes, 4, N) that contains the bounding boxes from the object detector
text_alpha (float): transparency label of the box wrapped around text labels.
ground_truth (bool): whether the prodived bounding boxes are ground-truth.
keyframe_idx (int): the index of keyframe in the clip.
repeat_frame (int): repeat each frame in draw_range for `repeat_frame` time for slow-motion effect.
"""
args = parse_args()
cfg = load_config(args)
scores_path = cfg.DEMO.SCORES_FILE_PATH
draw_static_boxes = cfg.DEMO.DRAW_STATIC_BOXES
assert repeat_frame >= 1, "`repeat_frame` must be a positive integer."
repeated_seq = range(0, len(frames))
repeated_seq = list(
itertools.chain.from_iterable(
itertools.repeat(x, repeat_frame) for x in repeated_seq))
frames, adjusted = self._adjust_frames_type(frames)
if keyframe_idx is None:
half_left = len(repeated_seq) // 2
half_right = (len(repeated_seq) + 1) // 2
else:
mid = int((keyframe_idx / len(frames)) * len(repeated_seq))
half_left = mid
half_right = len(repeated_seq) - mid
alpha_ls = np.concatenate([
np.linspace(0, 1, num=half_left),
np.linspace(1, 0, num=half_right),
])
text_alpha = text_alpha
frames = frames[repeated_seq]
img_ls = []
if mid is not None:
mid_frame = frames[mid]
else:
mid_frame = frames[half_left]
ids_boxes = [
self.tracker.advance(nbboxes[i], frame)
for i, frame in enumerate(frames)
]
mid_ids, _ = ids_boxes[mid if mid is not None else half_left]
with open(scores_path, 'a+') as f:
for i, (box_ids, bboxes) in enumerate(ids_boxes):
for box_id, bbox, pred in zip(box_ids, bboxes, preds):
f.write(
f'{self.tracker.current_task_id},{i},{box_id},{str(bbox.tolist()).strip("[]")},{str(pred.tolist()).strip("[]")}\n'
)
for i, (alpha, frame) in enumerate(zip(alpha_ls, frames)):
box_ids, frame_boxes = ids_boxes[i]
if not draw_static_boxes:
frame_preds = [
preds[mid_ids.index(box_id)]
if box_id in mid_ids else torch.zeros_like(preds[0])
for box_id in box_ids
]
draw_img = self.draw_one_frame(
frame,
preds=frame_preds,
bboxes=frame_boxes,
box_ids=box_ids,
alpha=1,
text_alpha=text_alpha,
)
else:
draw_img = self.draw_one_frame(
frame,
bboxes=frame_boxes,
box_ids=box_ids,
alpha=1,
text_alpha=text_alpha,
)
draw_img = self.draw_one_frame(
draw_img,
preds=preds,
bboxes=bboxes,
box_ids=mid_ids,
alpha=alpha,
text_alpha=text_alpha,
ground_truth=ground_truth,
)
if adjusted:
draw_img = draw_img.astype("float32") / 255
img_ls.append(draw_img)
return img_ls