xpra icon
Bug tracker and wiki

Ticket #350: python-frame-tracking.patch

File python-frame-tracking.patch, 11.2 KB (added by Antoine Martin, 6 years ago)

implements frame tracking in python

  • xpra/codecs/dec_avcodec/decoder.pyx

     
    3232ctypedef long AVPixelFormat
    3333
    3434cdef extern from "libavcodec/avcodec.h":
    35     ctypedef struct AVCodecContext:
    36         int width
    37         int height
    38         AVPixelFormat pix_fmt
    3935    ctypedef struct AVFrame:
    4036        uint8_t **data
    4137        int *linesize
     
    5046        uint8_t *data
    5147        int      size
    5248
     49    ctypedef struct AVCodecContext:
     50        int width
     51        int height
     52        AVPixelFormat pix_fmt
     53        int (*get_buffer)(AVCodecContext *c, AVFrame *pic)
     54        void (*release_buffer)(AVCodecContext *avctx, AVFrame *frame)
     55
    5356    AVPixelFormat PIX_FMT_NONE
    5457    AVCodecID CODEC_ID_H264
    5558
     59    #init and free:
    5660    void avcodec_register_all()
    5761    AVCodec *avcodec_find_decoder(AVCodecID id)
    5862    AVCodecContext *avcodec_alloc_context3(const AVCodec *codec)
     
    6266    int avcodec_close(AVCodecContext *avctx)
    6367    void av_free(void *ptr)
    6468
     69    #actual decoding:
    6570    void av_init_packet(AVPacket *pkt) nogil
    6671    void avcodec_get_frame_defaults(AVFrame *frame) nogil
    6772    int avcodec_decode_video2(AVCodecContext *avctx, AVFrame *picture,
    6873                                int *got_picture_ptr, const AVPacket *avpkt) nogil
    6974
     75    #buffer management:
     76    int avcodec_default_get_buffer(AVCodecContext *s, AVFrame *pic)
     77    void avcodec_default_release_buffer(AVCodecContext *s, AVFrame *pic)
     78
     79
    7080cdef extern from "dec_avcodec.h":
    7181    char *get_avcodec_version()
    7282    char **get_supported_colorspaces()
     
    104114    return get_avcodec_version()
    105115
    106116
     117#maps AVCodecContext to the Decoder that manages it
     118DECODERS = {}
     119
     120
     121cdef int avcodec_get_buffer(AVCodecContext *avctx, AVFrame *frame) with gil:
     122    log.info("avcodec_get_buffer(%s, %s)", hex(<unsigned long> avctx), hex(<unsigned long>  frame))
     123    return avcodec_default_get_buffer(avctx, frame)
     124
     125cdef void avcodec_release_buffer(AVCodecContext *avctx, AVFrame *frame) with gil:
     126    cdef unsigned long ctx_key = <unsigned long> avctx
     127    decoder = DECODERS.get(ctx_key)
     128    log.info("avcodec_release_buffer(%s, %s) decoder=%s", hex(<unsigned long> avctx), hex(<unsigned long>  frame), decoder)
     129    cdef unsigned long frame_key = <unsigned long> frame
     130    decoder.av_free(frame_key)
     131    #avcodec_default_release_buffer(avctx, frame)
     132
     133
     134
     135cdef class AVFrameWrapper:
     136    """
     137        Wraps an AVFrame so we can free it
     138        once both xpra and avcodec are done with it.
     139    """
     140    cdef AVCodecContext *avctx
     141    cdef AVFrame *frame
     142    cdef int av_freed
     143    cdef int xpra_freed
     144
     145    cdef set_context(self, AVCodecContext *avctx, AVFrame *frame):
     146        self.avctx = avctx
     147        self.frame = frame
     148        self.av_freed = 0
     149        self.xpra_freed = 0
     150
     151    def __dealloc__(self):
     152        #debug("CSCImage.__dealloc__()")
     153        assert self.av_freed, "AVFrameWrapper falling out of scope before being freed by avcodec!"
     154        assert self.xpra_freed, "AVFrameWrapper falling out of scope before being freed by xpra!"
     155        assert self.frame==NULL and self.avctx==NULL, "frame was freed by both, but not actually freed!"
     156
     157    def xpra_free(self):
     158        log.info("AVFrameWrapper.xpra_free()")
     159        self.xpra_freed = 1
     160        if self.av_freed==0:
     161            return False
     162        self.free()
     163        return True
     164
     165    def av_free(self):
     166        log.info("AVFrameWrapper.av_free()")
     167        self.av_freed = 1
     168        if self.xpra_freed==0:
     169            return False
     170        self.free()
     171        return True
     172
     173    def free(self):
     174        log.info("AVFrameWrapper.free()")
     175        if self.avctx!=NULL and self.frame!=NULL:
     176            avcodec_default_release_buffer(self.avctx, self.frame)
     177            #do we need this?
     178            #avcodec_free_frame(frame)
     179            self.avctx = NULL
     180            self.frame = NULL
     181
     182
     183class AVImageWrapper(ImageWrapper):
     184
     185    def free(self):                             #@DuplicatedSignature
     186        debug("CSCImageWrapper.free() av_frame=%s", self.av_frame)
     187        ImageWrapper.free(self)
     188        if self.av_frame:
     189            assert self.decoder, "no decoder set!"
     190            self.decoder.xpra_free(self.av_frame)
     191            self.av_frame = None
     192
     193
    107194cdef class Decoder:
    108195    cdef AVCodec *codec
    109196    cdef AVCodecContext *codec_ctx
    110197    cdef AVPixelFormat pix_fmt
    111198    cdef AVPixelFormat actual_pix_fmt
    112     cdef AVFrame *frame
    113199    cdef char *colorspace
    114     cdef object last_image
     200    cdef object avframes
    115201
    116202    def init_context(self, int width, int height, colorspace):
    117203        assert colorspace in COLORSPACES, "invalid colorspace: %s" % colorspace
     
    145231        self.codec_ctx.width = width
    146232        self.codec_ctx.height = height
    147233        self.codec_ctx.pix_fmt = self.pix_fmt
     234        self.codec_ctx.get_buffer = avcodec_get_buffer
     235        self.codec_ctx.release_buffer = avcodec_release_buffer
    148236        if avcodec_open2(self.codec_ctx, self.codec, NULL) < 0:
    149237            error("could not open codec")
    150238            self.clean_decoder()
    151239            return  False
    152 
    153         self.frame = avcodec_alloc_frame()
    154         if self.frame==NULL:
    155             error("could not allocate an AVFrame for decoding")
    156             self.clean_decoder()
    157             return  False
     240        #reference counting:
     241        self.avframes = {}
     242        global DECODERS
     243        cdef unsigned long ctx_key = <unsigned long> self.codec_ctx
     244        DECODERS[ctx_key] = self
    158245        return True
    159246
    160247    def clean(self):
    161         if self.last_image:
    162             #make sure the ImageWrapper does not reference memory
    163             #that is going to be freed!
    164             self.last_image.clone_pixel_data()
    165             self.last_image = None
    166248        self.clean_decoder()
    167249
    168250    def clean_decoder(self):
    169         if self.frame!=NULL:
    170             avcodec_free_frame(&self.frame)
    171             self.frame = NULL
     251        cdef unsigned long ctx_key = <unsigned long> self.codec_ctx
    172252        if self.codec_ctx!=NULL:
    173253            avcodec_close(self.codec_ctx)
    174254            av_free(self.codec_ctx)
     255            global DECODERS
     256            if ctx_key in DECODERS:
     257                DECODERS[ctx_key] = self
    175258            self.codec_ctx = NULL
    176259
    177260
     
    207290        cdef int len = 0
    208291        cdef int got_picture
    209292        cdef AVPacket avpkt
     293        cdef AVFrame *frame
    210294        assert self.codec_ctx!=NULL
    211295        assert self.codec!=NULL
    212         if self.last_image:
    213             #if another thread is still using this image
    214             #it is probably too late to prevent a race...
    215             #(it may be using the buffer directly by now)
    216             #but at least try to prevent new threads from
    217             #using the same buffer we are about to write to:
    218             self.last_image.clone_pixel_data()
    219             self.last_image = None
    220296        #copy input buffer into padded C buffer:
    221297        PyObject_AsReadBuffer(input, <const void**> &buf, &buf_len)
    222298        padded_buf = <unsigned char *> xmemalign(buf_len+128)
    223299        memcpy(padded_buf, buf, buf_len)
    224300        memset(padded_buf+buf_len, 0, 128)
    225301        #now safe to run without gil:
     302        frame = avcodec_alloc_frame()
     303        log.info("avcodec_alloc_frame()=%s", hex(<unsigned long> frame))
     304        if frame==NULL:
     305            raise Exception("avcodec could not allocate an AVFrame for decoding")
    226306        with nogil:
    227307            av_init_packet(&avpkt)
    228             avcodec_get_frame_defaults(self.frame)
    229308            avpkt.data = <uint8_t *> padded_buf
    230309            avpkt.size = buf_len
    231             len = avcodec_decode_video2(self.codec_ctx, self.frame, &got_picture, &avpkt)
     310            len = avcodec_decode_video2(self.codec_ctx, frame, &got_picture, &avpkt)
    232311            free(padded_buf)
    233312        if len < 0:
    234313            raise Exception("avcodec_decode_video2 failed to decode this frame")
    235314
    236315        #actual pixfmt:
    237         if self.pix_fmt!=self.frame.format:
    238             self.actual_pix_fmt = self.frame.format
     316        if self.pix_fmt!=frame.format:
     317            self.actual_pix_fmt = frame.format
    239318            debug("avcodec actual output pixel format is %s: %s" % (self.pix_fmt, self.get_actual_colorspace()))
    240319
    241320        #print("decompress image: colorspace=%s / %s" % (self.colorspace, self.get_colorspace()))
     
    254333                    height = (self.codec_ctx.height+1)>>1
    255334                else:
    256335                    raise Exception("invalid height divisor %s" % dy)
    257                 stride = self.frame.linesize[i]
     336                stride = frame.linesize[i]
    258337                size = height * stride
    259338                outsize += size
    260                 plane = PyBuffer_FromMemory(<void *>self.frame.data[i], size)
     339                plane = PyBuffer_FromMemory(<void *>frame.data[i], size)
    261340                out.append(plane)
    262341                strides.append(stride)
    263342        else:
    264             strides = self.frame.linesize[0]+self.frame.linesize[1]+self.frame.linesize[2]
     343            strides = frame.linesize[0]+frame.linesize[1]+frame.linesize[2]
    265344            outsize = self.codec_ctx.height * strides
    266             out = PyBuffer_FromMemory(<void *>self.frame.data[0], outsize)
     345            out = PyBuffer_FromMemory(<void *>frame.data[0], outsize)
    267346            nplanes = 0
    268347        if outsize==0:
    269348            raise Exception("output size is zero!")
    270349        debug("avcodec: %s bytes of %s", outsize, cs)
    271         img = ImageWrapper(0, 0, self.codec_ctx.width, self.codec_ctx.height, out, cs, 24, strides, nplanes)
    272         self.last_image = img
     350        img = AVImageWrapper(0, 0, self.codec_ctx.width, self.codec_ctx.height, out, cs, 24, strides, nplanes)
     351        img.decoder = self
     352        avframe_wrapper = AVFrameWrapper()
     353        avframe_wrapper.set_context(self.codec_ctx, frame)
     354        cdef unsigned long frame_key = <unsigned long> frame
     355        self.avframes[int(frame_key)] = avframe_wrapper
     356        log.info("avcodec adding %s with key=%s", avframe_wrapper, hex(frame_key))
     357        self.dump_frames()
     358        img.av_frame = avframe_wrapper
    273359        return img
    274360
     361    def dump_frames(self):
     362        log.info("frame keys=%s", [hex(x) for x in self.avframes.keys()])
     363       
     364
     365    def av_free(self, frame_key):
     366        log.info("av_free(%s)", hex(frame_key))
     367        avframe = self.avframes.get(int(frame_key))
     368        log.info("av_free(%s) avframe=%s", hex(frame_key), avframe)
     369        if avframe is None:
     370            self.dump_frames()
     371            raise Exception("frame not found for pointer %s" % (hex(frame_key)))
     372        if avframe.av_free():
     373            #frame has been freed - remove it from dict:
     374            del self.avframes[frame_key]
     375
     376    def xpra_free(self, avframe):
     377        log.info("xpra_free(%s)", avframe)
     378        if avframe.xpra_free():
     379            frame_key = None
     380            for key, frame in self.avframes.items():
     381                if frame==avframe:
     382                    frame_key = key
     383                    break
     384            log.info("xpra_free(%s) found key=%s", avframe, frame_key)
     385            assert frame_key
     386            #frame has been freed - remove it from dict:
     387            del self.avframes[frame_key]
     388
    275389    def get_colorspace(self):
    276390        return self.colorspace
    277391