if cfg.DEMO.THREAD_ENABLE:
frame_provider = ThreadVideoManager(cfg)
else:
frame_provider = VideoManager(cfg)
for task in tqdm.tqdm(run_demo(cfg, frame_provider)): # tqdm在长循环中添加一个进度提示信息 参数为一个迭代器
frame_provider.display(task)
cfg.DETECTION.ENABLE==False and cfg.DEMO.PREDS_BOXES == “” 进入else
cfg.DEMO.THREAD_ENABLE=False 创建一个VideoManager对象(frame_provider),其初始化中根据输入视频地址获取了视频参数并设置了输出结果视频的相关参数。其中宽、高、FPS采用输入视频的值(cfg中未给出),裁剪大小等使用cfg设置值。
run_demo为一个迭代器,yield task
首先设置种子、日志,打印日志
video_vis = VideoVisualizer(参数省略)
async_vis = AsyncVis(video_vis, n_workers=cfg.DEMO.NUM_VIS_INSTANCES)
if cfg.NUM_GPUS <= 1:
model = ActionPredictor(cfg=cfg, async_vis=async_vis)
else:
model = AsyncDemo(cfg=cfg, async_vis=async_vis)
seq_len = cfg.DATA.NUM_FRAMES * cfg.DATA.SAMPLING_RATE # 序列长度seq_len=帧数16*采样率5=80
回到run_demo函数体中:
设置序列长度seq_len=帧数16*采样率5=80
然后从frame_provider中获取task,依次跳入VideoManager的__iter__(self)方法和__next__(self)方法
for able_to_read, task in frame_provider:
class VideoManager:
def __iter__(self):
return self
def __next__(self):
self.id += 1 # 每次迭代id+1
task = TaskInfo()
task.img_height = self.display_height
......
frames = []
if len(self.buffer) != 0:
frames = self.buffer
was_read = True
while was_read and len(frames) < self.seq_length: #读入self.seq_length长度的帧到frames列表中,每帧为(720,1280,3)(读入测试视频大小)的ndarray
was_read, frame = self.cap.read() # self.cap为cv2.VideoCapture对象。
frames.append(frame)
if was_read and self.buffer_size != 0:
self.buffer = frames[-self.buffer_size :]
task.add_frames(self.id, frames)
task.num_buffer_frames = 0 if self.id == 0 else self.buffer_size
return was_read, task #若读取的帧不足seq_length,则was_read为False
再回到run_demo函数体中
for able_to_read, task in frame_provider:
if not able_to_read:
break
if task is None:
time.sleep(0.02)
continue
num_task += 1
model.put(task)
try:
task = model.get()
num_task -= 1
yield task
except IndexError:
continue
# 若在前面的每次循环中有没成功yield的,则再循环get,再yield
while num_task != 0:
......
循环中每次得到的task首先put进model,再从model中get,最后yield。若有没成功yield的,则再循环get,再yield。主要函数为model.put(task)和model.get()
当model为单GPU时的ActionPredictor对象
class ActionPredictor:
def put(self, task):
"""
Make prediction and put the results in `async_vis` task queue.
"""
task = self.predictor(task) # 处理输入数据(每个task的帧序列),输入到模型,得到X类的得分值
self.async_vis.get_indices_ls.append(task.id)
self.async_vis.put(task) # 加入到async_vis的task_queue(多进程队列)中
首先调用predictor.py中的Predictor对象的__call__方法, 处理输入数据(每个task的帧序列),输入到模型得到X类的得分值。另外将task.id和task分别加入到self.async_vis中。
下面细看predictor对象(Predictor类对象)
class Predictor:
def __init__(self, cfg, gpu_id=None):
if cfg.NUM_GPUS:
self.gpu_id = (
torch.cuda.current_device() if gpu_id is None else gpu_id
)
# Build the video model and print model statistics.
self.model = build_model(cfg, gpu_id=gpu_id)
self.model.eval()
self.cfg = cfg
if cfg.DETECTION.ENABLE:
self.object_detector = Detectron2Predictor(cfg, gpu_id=self.gpu_id)
logger.info("Start loading model weights.")
cu.load_test_checkpoint(cfg, self.model)
logger.info("Finish loading model weights")
其中主要有build_model和load_test_checkpoint两个重要函数
build_model函数在slowfast/models/build.py/build_model
def build_model(cfg, gpu_id=None):
......
# Construct the model
name = cfg.MODEL.MODEL_NAME
model = MODEL_REGISTRY.get(name)(cfg)
......
return model
前后省略部分为关于gpu的判断和设置
model = MODEL_REGISTRY.get(name)(cfg) 根据cfg.MODEL.MODEL_NAME构建相应模型类的对象 类定义在slowfast/models/video_model_builder.py中 类名只有SlowFast、ResNet、X3D。 在构建过程中会根据cfg.MODEL.ARCH(包含c2d、c2d_nopool、i3d、i3d_nopool、slow、slowfast、x3d) 设定核的尺寸。
from fvcore.common.registry import Registry
XXX_REGISTRY = Registry("XXX")
XXX_REGISTRY.__doc__ = """
#然后在类定义的上一行添加“@XXX_REGISTRY.register()”
@XXX_REGISTRY.register()
class X3D(nn.Module):
pass
@XXX_REGISTRY.register()
class ResNet(nn.Module)
pass
# 通过XXX_REGISTRY.get(name)方法,返回相应的类(的构造函数)
model=XXX_REGISTRY.get('X3D')('''类构造函数参数''')
加载模型参数:
首先检查cfg.TEST.CHECKPOINT_FILE_PATH,若为"",即没给出测试用checkpoint的地址, 则检查cfg.OUTPUT_DIR路径下的checkpoints文件夹中有无checkpoint(即自己训练产生的)。 若仍无,则检查cfg.TRAIN.CHECKPOINT_FILE_PATH
class Predictor:
def __call__(self, task):
......
frames, bboxes = task.frames, task.bboxes
......
# 以TEST_CROP_SIZE:356作为短边长度,将每帧数据缩放
# frames从(720,1280,3)的数组列表变为(356,632,3)的数组列表
frames = [
cv2_transform.scale(self.cfg.DATA.TEST_CROP_SIZE, frame) # TEST_CROP_SIZE:356
for frame in frames
]
# 数据预处理:归一化、标准化、调整维度、选择cfg.DATA.NUM_FRAMES帧、放入列表、增加维度
inputs = process_cv2_inputs(frames, self.cfg) # [(1,3,16,356,632)]
......
preds = self.model(inputs, bboxes) # 得到该片段 属于各类别的得分值
.....
task.add_action_preds(preds) # 将预测得分作为task的一项属性值
.....
return task
先提取出task中的数据做了预处理,包括从中采样出cfg.DATA.NUM_FRAMES个帧
因为task中总帧数为cfg.DATA.NUM_FRAMES * cfg.DATA.SAMPLING_RATE,所以采样cfg.DATA.NUM_FRAMES个帧,采样率即为cfg.DATA.SAMPLING_RATE。
经过模型得到该片段属于各类别的得分值,并将得分添加为task的一项属性值。
当model为多GPU时的AsyncDemo对象,最终同样是调用Predictor对象的__call__方法
pass
for task in tqdm.tqdm(run_demo(cfg, frame_provider)): # tqdm在长循环中添加一个进度提示信息 参数为一个迭代器
frame_provider.display(task)
回到demo,将run_demo()返回(yield)的一个task作为参数, 调用frame_provider.display(task),将task的frame写到指定文件(cfg.DEMO.OUTPUT_FILE)
为每一帧添加得分框的过程在Asyncis的_VisWorker的run中 调用draw_predictions -> video_vis.draw_clip_range -> self.draw_clip -> self.draw_one_frame :根据cfg.DEMO.VIS_MODE(“top-k"或"thres”)选择类别。 由于defaults.py中默认DEMO.VIS_MODE = “thres”, 而得分低于阈值,所以没有选出类别
在yaml中的DEMO下添加VIS_MODE: top-k