av/codec/mjpeg/extract.go

/*
DESCRIPTION
  extract.go provides an Extractor to get JPEG from RTP.

AUTHOR
  Saxon Nelson-Milton <saxon@ausocean.org>

LICENSE
  Copyright (C) 2017 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
*/

package mjpeg

import (
	"bytes"
	"fmt"
	"io"
	"time"
	"errors"

	"bitbucket.org/ausocean/av/protocol/rtp"
)

// Buffer sizes.
const (
	maxRTPSize = 1500 // Max ethernet transmission unit in bytes.
)

var (
	errNoQTable = errors.New("no quantization table")
	errReservedQ = errors.New("q value is reserved")
)

var defaultQuantisers = []byte{
        // Luma table.
        16,  11,  12,  14,  12,  10,  16,  14,
        13,  14,  18,  17,  16,  19,  24,  40,
        26,  24,  22,  22,  24,  49,  35,  37,
        29,  40,  58,  51,  61,  60,  57,  51,
        56,  55,  64,  72,  92,  78,  64,  68,
        87,  69,  55,  56,  80,  109, 81,  87,
        95,  98,  103, 104, 103, 62,  77,  113,
        121, 112, 100, 120, 92,  101, 103, 99,

        /* chroma table */
        17,  18,  18,  24,  21,  24,  47,  26,
        26,  47,  99,  66,  56,  66,  99,  99,
        99,  99,  99,  99,  99,  99,  99,  99,
        99,  99,  99,  99,  99,  99,  99,  99,
        99,  99,  99,  99,  99,  99,  99,  99,
        99,  99,  99,  99,  99,  99,  99,  99,
        99,  99,  99,  99,  99,  99,  99,  99,
        99,  99,  99,  99,  99,  99,  99,  99,

}

// Extractor is an Extractor for extracting JPEG from an RTP stream.
type Extractor struct {
	buf *bytes.Buffer // Holds the current JPEG image.
	dst io.Writer     // The destination we'll be writing extracted NALUs to.
}

// NewExtractor returns a new Extractor.
func NewExtractor() *Extractor { return &Extractor{} }

// Extract will continously read RTP packets from src containing JPEG (in RTP
// payload format) and extract the JPEG images, sending them to dst. This
// function expects that each read from src will provide a single RTP packet.
func (e *Extractor) Extract(dst io.Writer, src io.Reader, delay time.Duration) error {
	buf := make([]byte, maxRTPSize)

	var (
		qTables [128][128]byte
		qTablesLen [128]byte
	)

	for {
		n, err := src.Read(buf)
		switch err {
		case nil: // Do nothing.
		case io.EOF:
			return nil
		default:
			return fmt.Errorf("source read error: %v\n", err)
		}

		// Get payload from RTP packet.
		p, err := rtp.Payload(buf[:n])
		if err != nil {
			return fmt.Errorf("could not get RTP payload, failed with err: %v\n", err)
		}

		b := newByteStream(p)
		_ = b.get8() // Ignore type-specific flag

		var (
			off    = b.get24() // Fragment offset.
			t      = b.get8()  // Type.
			q      = b.get8()  // Quantization value.
			width  = b.get8()  // Picture width.
			height = b.get8()  // Picture height.
			dri    int         // Restart interval.
		)

		if t&0x40 != 0 {
			dri = b.get16()
			_ = b.get16() // Ignore restart count.
			t &= ^0x40
		}

		if t > 1 {
			panic("unimplemented RTP/JPEG type")
		}

		// Parse quantization table if our offset is 0.
		if off == 0 {
			var qTable []byte
			var qLen int

			if q > 127 {
				_ = b.get8() // Ignore first byte (reserved for future use).
				prec := b.get8() // The size of coefficients.
				qLen = b.get16() // The length of the quantization table.

				if prec != 0 {
					panic("unsupported precision")
				}

				if qLen > 0 {
					qTable = b.getBuf(qLen)

					if q < 255 {
						if qTablesLen[q-128] == 0 && qLen <= 128 {
							copy(qTables[q-128][:],qTable)
							qTablesLen[q-128] = byte(qLen)
						}
					}
				} else {
					if q == 255 {
						return errNoQTable
					}

					if qTablesLen[q-128] == 0 {
						return fmt.Errorf("no quantization tables known for q %d yet",q)
					}

					qTable = qTables[q-128][:]
					qLen = int(qTablesLen[q-128])
				}
			} else { // q <= 127
				if q == 0 || q > 99 {
					return errReservedQ
				}
				qTable = defaultQTable(q)
				qLen = len(qTable)
			}

			e.buf.Reset()

			err = writeHeader(e.buf, t, width, height, qLen / 64, dri, qTable)
			if err != nil {
				return fmt.Errorf("could not write JPEG header: %w",err)
			}
		}

		if e.buf.Len() == 0 {
			// Must have missed start of frame? So ignore and wait for start.
			continue
		}

		// TODO: check that timestamp is consistent
		// This will need expansion to RTP package to create Timestamp parsing func.

		// TODO: could also check offset with how many bytes we currently have
		// to determine if there are missing frames.

		// Write frame data
		err = b.writeTo(e.buf,b.remaining())
		if err != nil {
			return fmt.Errorf("could not write remaining frame data to output buffer: %w",err)
		}

		m, err := rtp.Marker(buf[:n])
		if err != nil {
			return fmt.Errorf("could not read RTP marker: %w",err)
		}

		if m {
			_,err = e.buf.Write([]byte{0xff,codeEOI})
			if err != nil {
				return fmt.Errorf("could not write EOI marker: %w",err)
			}

			_,err = e.buf.WriteTo(dst)
			if err != nil {
				return fmt.Errorf("could not write JPEG to dst: %w",err)
			}
		}
	}
}

type byteStream struct {
	bytes []byte
	i     int
}

func newByteStream(b []byte) *byteStream { return &byteStream{bytes: b} }

func (b *byteStream) get24() int {
	v := int(b.bytes[b.i])<<16 | int(b.bytes[b.i+1])<<8 | int(b.bytes[b.i+2])
	b.i += 3
	return v
}

func (b *byteStream) get8() int {
	v := int(b.bytes[b.i])
	b.i++
	return v
}

func (b *byteStream) get16() int {
	v :=  int(b.bytes[b.i])<<8 | int(b.bytes[b.i+1])
	b.i += 2
	return v
}

func (b *byteStream) getBuf(n int) []byte {
	v := b.bytes[b.i:b.i+n]
	b.i += n
	return v
}

func (b *byteStream) remaining() int {
	return len(b.bytes) - b.i
}

func (b *byteStream) writeTo(w io.Writer, n int) error {
	_n,err := w.Write(b.bytes[b.i:n])
	b.i += _n
	if err != nil {
		return err
	}
	return nil
}

func defaultQTable(q int) []byte {
	f := clip(q,q,99)
	const tabLen = 128
	tab := make([]byte,tabLen)

	if q < 50 {
		q = 5000 / f
	} else {
		q = 200 - f*2
	}

	for i := 0; i < tabLen; i++ {
		v := (int(defaultQuantisers[i])*q + 50) / 100
		v = clip(v,1,255)
		tab[i] = byte(v)
	}
	return tab
}

func clip(v, min, max int) int {
	if v < min {
		return min
	}

	if v > max {
		return max
	}

	return v
}