def next()

in python/vmaf/tools/reader.py [0:0]

• 52 lines of code
• 23 McCabe index (conditional complexity)

    def next(self, format='uint'):

        assert format == 'uint' or format == 'float'

        y_width = self.width
        y_height = self.height
        uv_w_multiplier, uv_h_multiplier = self._get_uv_width_height_multiplier()
        uv_width = int(y_width * uv_w_multiplier)
        uv_height = int(y_height * uv_h_multiplier)

        if self._is_8bit():
            pix_type = np.uint8
            word = 1
        elif self._is_10bitle() or self._is_12bitle() or self._is_16bitle():
            pix_type = np.uint16
            word = 2
        else:
            assert False

        y = np.frombuffer(self.file.read(y_width * y_height * word), pix_type)

        if y.size == 0:
            raise StopIteration

        if uv_width == 0 and uv_height == 0:
            u = None
            v = None
        elif uv_width > 0 and uv_height > 0:
            u = np.frombuffer(self.file.read(uv_width * uv_height * word), pix_type)
            if u.size == 0:
                raise StopIteration
            v = np.frombuffer(self.file.read(uv_width * uv_height * word), pix_type)
            if v.size == 0:
                raise StopIteration
        else:
            assert False, f'Unsupported uv_width and uv_height: {uv_width}, {uv_height}'

        y = y.reshape(y_height, y_width)
        u = u.reshape(uv_height, uv_width) if u is not None else None
        v = v.reshape(uv_height, uv_width) if v is not None else None

        if format == 'uint':
            return y, u, v

        elif format == 'float':
            if self._is_8bit():
                bit_depth = 8
            elif self._is_10bitle():
                bit_depth = 10
            elif self._is_12bitle():
                bit_depth = 12
            elif self._is_16bitle():
                bit_depth = 16
            else:
                assert False

            y = self.convert_format(y, bit_depth)
            u = self.convert_format(u, bit_depth) if u is not None else None
            v = self.convert_format(v, bit_depth) if v is not None else None
            return y, u, v

        else:
            assert False