diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2cbd7d33..6ff23d1b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -20,6 +20,8 @@ jobs: with: go-version: ${{ matrix.go }} + - run: sudo apt install -y libavformat-dev libswscale-dev + - run: make test-nodocker - if: matrix.go == '1.16' diff --git a/Makefile b/Makefile index 87e29beb..d4282286 100644 --- a/Makefile +++ b/Makefile @@ -39,7 +39,7 @@ format: define DOCKERFILE_TEST FROM $(BASE_IMAGE) -RUN apk add --no-cache make docker-cli git gcc musl-dev +RUN apk add --no-cache make docker-cli git gcc musl-dev pkgconfig ffmpeg-dev WORKDIR /s COPY go.mod go.sum ./ RUN go mod download diff --git a/README.md b/README.md index 3be6380c..429856e0 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,7 @@ Features: * [client-read-options](examples/client-read-options/main.go) * [client-read-pause](examples/client-read-pause/main.go) * [client-read-h264](examples/client-read-h264/main.go) +* [client-read-h264-convert-to-jpeg](examples/client-read-h264-convert-to-jpeg/main.go) * [client-read-h264-save-to-disk](examples/client-read-h264-save-to-disk/main.go) * [client-read-aac](examples/client-read-aac/main.go) * [client-publish-h264](examples/client-publish-h264/main.go) diff --git a/examples/client-read-h264-convert-to-jpeg/h264decoder.go b/examples/client-read-h264-convert-to-jpeg/h264decoder.go new file mode 100644 index 00000000..f55bc7ee --- /dev/null +++ b/examples/client-read-h264-convert-to-jpeg/h264decoder.go @@ -0,0 +1,144 @@ +package main + +import ( + "fmt" + "image" + "unsafe" +) + +// #cgo pkg-config: libavcodec libavutil libswscale +// #include +// #include +// #include +import "C" + +func frameData(frame *C.AVFrame) **C.uint8_t { + return (**C.uint8_t)(unsafe.Pointer(&frame.data[0])) +} + +func frameLineSize(frame *C.AVFrame) *C.int { + return (*C.int)(unsafe.Pointer(&frame.linesize[0])) +} + +// h264Decoder is a wrapper around ffmpeg's H264 decoder. +type h264Decoder struct { + codecCtx *C.AVCodecContext + avPacket C.AVPacket + srcFrame *C.AVFrame + swsCtx *C.struct_SwsContext + dstFrame *C.AVFrame + dstFramePtr []uint8 +} + +// newH264Decoder allocates a new h264Decoder. +func newH264Decoder() (*h264Decoder, error) { + codec := C.avcodec_find_decoder(C.AV_CODEC_ID_H264) + if codec == nil { + return nil, fmt.Errorf("avcodec_find_decoder() failed") + } + + codecCtx := C.avcodec_alloc_context3(codec) + if codecCtx == nil { + return nil, fmt.Errorf("avcodec_alloc_context3() failed") + } + + res := C.avcodec_open2(codecCtx, codec, nil) + if res < 0 { + C.avcodec_close(codecCtx) + return nil, fmt.Errorf("avcodec_open2() failed") + } + + srcFrame := C.av_frame_alloc() + if srcFrame == nil { + C.avcodec_close(codecCtx) + return nil, fmt.Errorf("av_frame_alloc() failed") + } + + avPacket := C.AVPacket{} + C.av_init_packet(&avPacket) + + return &h264Decoder{ + codecCtx: codecCtx, + srcFrame: srcFrame, + avPacket: avPacket, + }, nil +} + +// close closes the decoder. +func (d *h264Decoder) close() { + if d.dstFrame != nil { + C.av_frame_free(&d.dstFrame) + } + + if d.swsCtx != nil { + C.sws_freeContext(d.swsCtx) + } + + C.av_frame_free(&d.srcFrame) + C.avcodec_close(d.codecCtx) +} + +func (d *h264Decoder) decode(nalu []byte) (image.Image, error) { + nalu = append([]uint8{0x00, 0x00, 0x00, 0x01}, []uint8(nalu)...) + + // send frame to decoder + d.avPacket.data = (*C.uint8_t)(C.CBytes(nalu)) + defer C.free(unsafe.Pointer(d.avPacket.data)) + d.avPacket.size = C.int(len(nalu)) + res := C.avcodec_send_packet(d.codecCtx, &d.avPacket) + if res < 0 { + return nil, nil + } + + // receive frame if available + res = C.avcodec_receive_frame(d.codecCtx, d.srcFrame) + if res < 0 { + return nil, nil + } + + // if frame size has changed, allocate needed objects + if d.dstFrame == nil || d.dstFrame.width != d.srcFrame.width || d.dstFrame.height != d.srcFrame.height { + if d.dstFrame != nil { + C.av_frame_free(&d.dstFrame) + } + + if d.swsCtx != nil { + C.sws_freeContext(d.swsCtx) + } + + d.swsCtx = C.sws_getContext(d.srcFrame.width, d.srcFrame.height, C.AV_PIX_FMT_YUV420P, // d.codecCtx.pix_fmt, + d.srcFrame.width, d.srcFrame.height, C.AV_PIX_FMT_RGBA, C.SWS_BILINEAR, nil, nil, nil) + if d.swsCtx == nil { + return nil, fmt.Errorf("sws_getContext() err") + } + + d.dstFrame = C.av_frame_alloc() + d.dstFrame.format = C.AV_PIX_FMT_RGBA + d.dstFrame.width = d.srcFrame.width + d.dstFrame.height = d.srcFrame.height + d.dstFrame.color_range = C.AVCOL_RANGE_JPEG + res = C.av_frame_get_buffer(d.dstFrame, 32) + if res < 0 { + return nil, fmt.Errorf("av_frame_get_buffer() err") + } + + dstFrameSize := C.av_image_get_buffer_size((int32)(d.dstFrame.format), d.dstFrame.width, d.dstFrame.height, 1) + d.dstFramePtr = (*[1 << 30]uint8)(unsafe.Pointer(d.dstFrame.data[0]))[:dstFrameSize:dstFrameSize] + } + + // convert frame from YUV420 to RGB + res = C.sws_scale(d.swsCtx, frameData(d.srcFrame), frameLineSize(d.srcFrame), + 0, d.codecCtx.height, frameData(d.dstFrame), frameLineSize(d.dstFrame)) + if res < 0 { + return nil, fmt.Errorf("sws_scale() err") + } + + // embed frame into an image.Image + return &image.RGBA{ + Pix: d.dstFramePtr, + Stride: 4 * (int)(d.dstFrame.width), + Rect: image.Rectangle{ + Max: image.Point{(int)(d.dstFrame.width), (int)(d.dstFrame.height)}, + }, + }, nil +} diff --git a/examples/client-read-h264-convert-to-jpeg/main.go b/examples/client-read-h264-convert-to-jpeg/main.go new file mode 100644 index 00000000..5718cca0 --- /dev/null +++ b/examples/client-read-h264-convert-to-jpeg/main.go @@ -0,0 +1,147 @@ +package main + +import ( + "image" + "image/jpeg" + "log" + "os" + "strconv" + "time" + + "github.com/aler9/gortsplib" + "github.com/aler9/gortsplib/pkg/base" + "github.com/aler9/gortsplib/pkg/rtph264" + "github.com/pion/rtp" +) + +// This example shows how to +// 1. connect to a RTSP server and read all tracks on a path +// 2. check whether there's a H264 track +// 3. decode the H264 track to raw frames +// 4. encode the frames into JPEG images and save them on disk +// This example requires the ffmpeg libraries, that can be installed in this way: +// apt install -y libavformat-dev libswscale-dev gcc pkg-config + +func saveToFile(img image.Image) error { + // create file + fname := strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10) + ".jpg" + f, err := os.Create(fname) + if err != nil { + panic(err) + } + defer f.Close() + + log.Println("saving", fname) + + // convert to jpeg + return jpeg.Encode(f, img, &jpeg.Options{ + Quality: 60, + }) +} + +func main() { + c := gortsplib.Client{} + + // parse URL + u, err := base.ParseURL("rtsp://localhost:8554/mystream") + if err != nil { + panic(err) + } + + // connect to the server + err = c.Start(u.Scheme, u.Host) + if err != nil { + panic(err) + } + defer c.Close() + + // get available methods + _, err = c.Options(u) + if err != nil { + panic(err) + } + + // find published tracks + tracks, baseURL, _, err := c.Describe(u) + if err != nil { + panic(err) + } + + // find the H264 track + h264Track := func() int { + for i, track := range tracks { + if _, ok := track.(*gortsplib.TrackH264); ok { + return i + } + } + return -1 + }() + if h264Track < 0 { + panic("H264 track not found") + } + + // setup RTP->H264 decoder + dec := rtph264.NewDecoder() + + // setup H264->raw frames decoder + h264dec, err := newH264Decoder() + if err != nil { + panic(err) + } + defer h264dec.close() + + // called when a RTP packet arrives + saveCount := 0 + c.OnPacketRTP = func(trackID int, payload []byte) { + if trackID != h264Track { + return + } + + // parse RTP packet + var pkt rtp.Packet + err := pkt.Unmarshal(payload) + if err != nil { + return + } + + // decode H264 NALUs from RTP packet + nalus, _, err := dec.Decode(&pkt) + if err != nil { + return + } + + // decode raw frames from H264 NALUs + for _, nalu := range nalus { + img, err := h264dec.decode(nalu) + if err != nil { + panic(err) + } + + // wait for a frame + if img == nil { + continue + } + + // convert frame to JPEG and save to file + err = saveToFile(img) + if err != nil { + panic(err) + } + + saveCount++ + if saveCount == 5 { + log.Printf("saved 5 images, exiting") + os.Exit(1) + } + } + } + + // start reading tracks + err = c.SetupAndPlay(tracks, baseURL) + if err != nil { + panic(err) + } + + // wait until a fatal error + panic(c.Wait()) +}