X3D代码理解之demo(cfg)

demo(cfg)

       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(cfg, frame_provider)

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

• 创建VideoVisualizer对象(video_vis)，根据cfg初始化属性值
• 创建AsyncVis对象(async_vis)，其成员主要包括两个mp.Queue()和包含cpu个数个_VisWorker对象（工作进程）的列表procs
• 若单GPU，创建ActionPredictor(model)，其成员包括一个Predictor对象(predictor)和AsyncVis对象(async_vis)；
• 若多GPU，创建AsyncDemo(model)，其成员包括一个AsycnActionPredictor对象和AsyncVis对象(async_vis)
  • AsycnActionPredictor对象的成员主要包括两个mp.Queue()和包含gpu个数个_Predictor对象的列表procs。(类似AsyncVis）
  • _Predictor对象的run方法中创建Predictor对象

回到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

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

• task为封装了一些数据成员的TaskInfo类，表示每次处理的一个视频剪辑。成员包括frames，id，bboxes，action_preds，num_buffer_frames，img_height，img_width，crop_size，clip_vis_size。
• 使用自身的一些成员值为task成员赋值
• 每次读取self.seq_length长度的帧到frames列表，用作输入，更新task.frames和task.num_buffer_frames
• 返回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.put(task)

当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类对象）

Predictor.init(self, cfg, gpu_id=None)

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

build_model(cfg, gpu_id=None)

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） 设定核的尺寸。

Registry用法小结（在不确定选择哪个类时使用）

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')('''类构造函数参数''')

cu.load_test_checkpoint(cfg, self.model)

加载模型参数：
首先检查cfg.TEST.CHECKPOINT_FILE_PATH，若为""，即没给出测试用checkpoint的地址， 则检查cfg.OUTPUT_DIR路径下的checkpoints文件夹中有无checkpoint（即自己训练产生的）。 若仍无，则检查cfg.TRAIN.CHECKPOINT_FILE_PATH

Predictor.call(self, task)

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__方法

model.get(task)

pass

demo(cfg):

        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