av/codec/mjpeg/extract.go

289 lines
6.6 KiB
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
}