コード例 #1
0
ファイル: cutter.py プロジェクト: Hanlen520/stagesepx
    def compare_frame_list(
            self, src: typing.List[np.ndarray],
            target: typing.List[np.ndarray]) -> typing.List[float]:
        """
        core method about how to compare two lists of ndarray and get their ssim/mse/psnr
        you can overwrite this method to implement your own algo
        see https://github.com/williamfzc/stagesepx/issues/136

        :param src:
        :param target:
        :return:
        """
        # find the min ssim and the max mse / psnr
        ssim = 1.0
        mse = 0.0
        psnr = 0.0

        for part_index, (each_start, each_end) in enumerate(zip(src, target)):
            part_ssim = toolbox.compare_ssim(each_start, each_end)
            if part_ssim < ssim:
                ssim = part_ssim

            # mse is very sensitive
            part_mse = toolbox.calc_mse(each_start, each_end)
            if part_mse > mse:
                mse = part_mse

            part_psnr = toolbox.calc_psnr(each_start, each_end)
            if part_psnr > psnr:
                psnr = part_psnr
            logger.debug(
                f"part {part_index}: ssim={part_ssim}; mse={part_mse}; psnr={part_psnr}"
            )
        return [ssim, mse, psnr]
コード例 #2
0
ファイル: cutter.py プロジェクト: zsg555/stagesepx
        def _compare_frame_list(
                src: typing.List[np.ndarray],
                target: typing.List[np.ndarray]) -> typing.List[float]:
            # find the min ssim and the max mse / psnr
            ssim = 1.0
            mse = 0.0
            psnr = 0.0

            for part_index, (each_start,
                             each_end) in enumerate(zip(src, target)):
                part_ssim = toolbox.compare_ssim(each_start, each_end)
                if part_ssim < ssim:
                    ssim = part_ssim

                # mse is very sensitive
                part_mse = toolbox.calc_mse(each_start, each_end)
                if part_mse > mse:
                    mse = part_mse

                part_psnr = toolbox.calc_psnr(each_start, each_end)
                if part_psnr > psnr:
                    psnr = part_psnr
                logger.debug(
                    f"part {part_index}: ssim={part_ssim}; mse={part_mse}; psnr={part_psnr}"
                )
            return [ssim, mse, psnr]
コード例 #3
0
ファイル: cutter.py プロジェクト: lshvisual/stagesepx
    def _convert_video_into_range_list(self,
                                       video: VideoObject,
                                       block: int = None,
                                       *args,
                                       **kwargs) -> typing.List[VideoCutRange]:
        if not block:
            block = 2

        range_list: typing.List[VideoCutRange] = list()
        with toolbox.video_capture(video.path) as cap:
            logger.debug(
                f'total frame count: {video.frame_count}, size: {video.frame_size}'
            )

            # load the first two frames
            _, start = cap.read()
            start_frame_id = toolbox.get_current_frame_id(cap)
            start_frame_time = toolbox.get_current_frame_time(cap)

            toolbox.video_jump(cap, self.step + 1)
            ret, end = cap.read()
            end_frame_id = toolbox.get_current_frame_id(cap)
            end_frame_time = toolbox.get_current_frame_time(cap)

            # hook
            start = self._apply_hook(start_frame_id, start)

            # check block
            if not self.is_block_valid(start, block):
                logger.warning(
                    'array split does not result in an equal division, set block to 1'
                )
                block = 1

            while ret:
                # hook
                end = self._apply_hook(end_frame_id, end, *args, **kwargs)

                logger.debug(
                    f'computing {start_frame_id}({start_frame_time}) & {end_frame_id}({end_frame_time}) ...'
                )
                start_part_list = self.pic_split(start, block)
                end_part_list = self.pic_split(end, block)

                # find the min ssim and the max mse / psnr
                ssim = 1.
                mse = 0.
                psnr = 0.
                for part_index, (each_start, each_end) in enumerate(
                        zip(start_part_list, end_part_list)):
                    part_ssim = toolbox.compare_ssim(each_start, each_end)
                    if part_ssim < ssim:
                        ssim = part_ssim

                    # mse is very sensitive
                    part_mse = toolbox.calc_mse(each_start, each_end)
                    if part_mse > mse:
                        mse = part_mse

                    part_psnr = toolbox.calc_psnr(each_start, each_end)
                    if part_psnr > psnr:
                        psnr = part_psnr
                    logger.debug(
                        f'part {part_index}: ssim={part_ssim}; mse={part_mse}; psnr={part_psnr}'
                    )
                logger.debug(
                    f'between {start_frame_id} & {end_frame_id}: ssim={ssim}; mse={mse}; psnr={psnr}'
                )

                range_list.append(
                    VideoCutRange(
                        video,
                        start=start_frame_id,
                        end=end_frame_id,
                        ssim=[ssim],
                        mse=[mse],
                        psnr=[psnr],
                        start_time=start_frame_time,
                        end_time=end_frame_time,
                    ))

                # load the next one
                start = end
                start_frame_id, end_frame_id = end_frame_id, end_frame_id + self.step
                start_frame_time = end_frame_time

                toolbox.video_jump(cap, end_frame_id)
                ret, end = cap.read()
                end_frame_time = toolbox.get_current_frame_time(cap)

        return range_list
コード例 #4
0
ファイル: cutter.py プロジェクト: liangjz/stagesepx
    def _convert_video_into_range_list(self,
                                       video: VideoObject,
                                       block: int = None,
                                       *args,
                                       **kwargs) -> typing.List[VideoCutRange]:

        range_list: typing.List[VideoCutRange] = list()
        logger.info(
            f"total frame count: {video.frame_count}, size: {video.frame_size}"
        )

        # load the first two frames
        video_operator = video.get_operator()
        cur_frame = video_operator.get_frame_by_id(1)
        next_frame = video_operator.get_frame_by_id(1 + self.step)

        # hook
        cur_frame.data = self._apply_hook(cur_frame.frame_id, cur_frame.data)

        # check block
        if not block:
            block = 2
        if not self.is_block_valid(cur_frame.data, block):
            logger.warning(
                "array split does not result in an equal division, set block to 1"
            )
            block = 1

        while True:
            # hook
            next_frame.data = self._apply_hook(next_frame.frame_id,
                                               next_frame.data, *args,
                                               **kwargs)

            logger.debug(
                f"computing {cur_frame.frame_id}({cur_frame.timestamp}) & {next_frame.frame_id}({next_frame.timestamp}) ..."
            )
            start_part_list = self.pic_split(cur_frame.data, block)
            end_part_list = self.pic_split(next_frame.data, block)

            # find the min ssim and the max mse / psnr
            ssim = 1.0
            mse = 0.0
            psnr = 0.0
            for part_index, (each_start, each_end) in enumerate(
                    zip(start_part_list, end_part_list)):
                part_ssim = toolbox.compare_ssim(each_start, each_end)
                if part_ssim < ssim:
                    ssim = part_ssim

                # mse is very sensitive
                part_mse = toolbox.calc_mse(each_start, each_end)
                if part_mse > mse:
                    mse = part_mse

                part_psnr = toolbox.calc_psnr(each_start, each_end)
                if part_psnr > psnr:
                    psnr = part_psnr
                logger.debug(
                    f"part {part_index}: ssim={part_ssim}; mse={part_mse}; psnr={part_psnr}"
                )
            logger.debug(
                f"between {cur_frame.frame_id} & {next_frame.frame_id}: ssim={ssim}; mse={mse}; psnr={psnr}"
            )

            range_list.append(
                VideoCutRange(
                    video,
                    start=cur_frame.frame_id,
                    end=next_frame.frame_id,
                    ssim=[ssim],
                    mse=[mse],
                    psnr=[psnr],
                    start_time=cur_frame.timestamp,
                    end_time=next_frame.timestamp,
                ))

            # load the next one
            cur_frame = next_frame
            next_frame = video_operator.get_frame_by_id(next_frame.frame_id +
                                                        self.step)
            if next_frame is None:
                break

        return range_list