webcam/webcam.go

package webcam

import (
	"bytes"
	"errors"
	"fmt"
	"image"
	"image/jpeg"
	"image/png"
	"os"
	"sort"
	"syscall"
	"time"
	"unsafe"

	"github.com/creack/goselect"
)

type Webcam struct {
	fd      uintptr
	file    *os.File
	path    string
	timeout time.Duration
	buffers [][]byte
}

func Open(path string, timeout time.Duration) (*Webcam, error) {
	f, err := os.OpenFile(path, syscall.O_RDWR|syscall.O_NONBLOCK, 0)
	if err != nil {
		return nil, err
	}

	w := new(Webcam)
	w.fd = f.Fd()
	w.file = f
	w.path = path
	w.timeout = timeout

	v, s, err := w.checkCapabilities()
	if err != nil {
		return nil, err
	}
	if v {
		return nil, errors.New("Not a video capture device")
	}
	if s {
		return nil, errors.New("Device does not support the streaming I/O method")
	}
	return w, nil
}

type requestbuffers struct {
	count    uint32
	t        uint32
	memory   uint32
	reserved [2]uint32
}

type buffer struct {
	index     uint32
	t         uint32
	bytesused uint32
	flags     uint32
	field     uint32
	timestamp syscall.Timeval
	timecode  struct {
		t        uint32
		flags    uint32
		frames   uint8
		seconds  uint8
		minutes  uint8
		hours    uint8
		userbits [4]uint8
	}
	sequence uint32
	memory   uint32
	/*
		union {
		        __u32           offset;
		        unsigned long   userptr;
		        struct v4l2_plane *planes;
		        __s32           fd;
		} m;
	*/
	offset    int64 // m
	length    uint32
	reserved2 uint32
	reserved  uint32
}

type format struct {
	t        uint32
	pix      pix
	raw_data [155]byte // placeholder
}
type pix struct {
	empty        uint32 // placeholder
	width        uint32
	height       uint32
	pixelformat  uint32
	field        uint32
	bytesperline uint32
	sizeimage    uint32
	colorspace   uint32
	priv         uint32
	flags        uint32
	ycbcr_enc    uint32
	quantization uint32
	// commented out for 'empty' placeholder
	//xfer_func    uint32
}

type capability struct {
	driver       [16]byte
	card         [32]byte
	bus_info     [32]byte
	version      uint32
	capabilities uint32
	device_caps  uint32
	reserved     [3]uint32
}

type fmtdesc struct {
	index       uint32
	t           uint32
	flags       uint32
	description [32]byte
	pixelformat uint32
	reserved    [4]uint32
}

type frmsizeenum struct {
	index        uint32
	pixel_format uint32
	t            uint32
	/*
		struct v4l2_frmsize_stepwise {
			__u32 min_width
			__u32 max_width
			__u32 step_width
			__u32 min_height
			__u32 max_height
			__u32 step_height
		};
		struct v4l2_frmsize_discrete {
			__u32 width
			__u32 height
		};
	*/
	frmsize  [6]uint32
	reserved [2]uint32
}

const (
	_IOC_READ  = 2
	_IOC_WRITE = 1

	_IOC_NRBITS    = 8
	_IOC_TYPEBITS  = 8
	_IOC_SIZEBITS  = 14
	_IOC_NRSHIFT   = 0
	_IOC_TYPESHIFT = _IOC_NRSHIFT + _IOC_NRBITS
	_IOC_SIZESHIFT = _IOC_TYPESHIFT + _IOC_TYPEBITS
	_IOC_DIRSHIFT  = _IOC_SIZESHIFT + _IOC_SIZEBITS

	_V4L2_BUF_TYPE_VIDEO_CAPTURE = 1
	_V4L2_MEMORY_MMAP            = 1
	_V4L2_FIELD_ANY              = 0

	_V4L2_CAP_STREAMING     = 0x04000000
	_V4L2_CAP_VIDEO_CAPTURE = 0x00000001

	_V4L2_FRMSIZE_TYPE_DISCRETE   = 1
	_V4L2_FRMSIZE_TYPE_CONTINUOUS = 2
	_V4L2_FRMSIZE_TYPE_STEPWISE   = 3
)

var (
	_VIDIOC_QUERYCAP        = _IOR('V', 0, unsafe.Sizeof(capability{}))
	_VIDIOC_ENUM_FMT        = _IOWR('V', 2, unsafe.Sizeof(fmtdesc{}))
	_VIDIOC_S_FMT           = _IOWR('V', 5, unsafe.Sizeof(format{}))
	_VIDIOC_REQBUFS         = _IOWR('V', 8, unsafe.Sizeof(requestbuffers{}))
	_VIDIOC_QUERYBUF        = _IOWR('V', 9, unsafe.Sizeof(buffer{}))
	_VIDIOC_QBUF            = _IOWR('V', 15, unsafe.Sizeof(buffer{}))
	_VIDIOC_DQBUF           = _IOWR('V', 17, unsafe.Sizeof(buffer{}))
	_VIDIOC_STREAMON        = _IOW('V', 18, unsafe.Sizeof(int32(0)))
	_VIDIOC_ENUM_FRAMESIZES = _IOWR('V', 74, unsafe.Sizeof(frmsizeenum{}))
)

func _IOC(dir, t, nr, size uintptr) uintptr {
	return dir<<_IOC_DIRSHIFT |
		t<<_IOC_TYPESHIFT |
		nr<<_IOC_NRSHIFT |
		size<<_IOC_SIZESHIFT
}

func _IOWR(t rune, i, size uintptr) uintptr {
	return _IOC(_IOC_READ|_IOC_WRITE, uintptr(t), i, size)
}

func _IOW(t rune, i, size uintptr) uintptr {
	return _IOC(_IOC_WRITE, uintptr(t), i, size)
}

func _IOR(t rune, i, size uintptr) uintptr {
	return _IOC(_IOC_READ, uintptr(t), i, size)
}

func ioctl(fd, cmd, ptr uintptr) error {
	_, _, ep := syscall.Syscall(syscall.SYS_IOCTL, fd, cmd, ptr)
	if ep != 0 {
		return syscall.Errno(ep)
	}
	return nil
}

func (w *Webcam) checkCapabilities() (video, stream bool, err error) {
	c := new(capability)
	err = ioctl(w.fd, _VIDIOC_QUERYCAP, uintptr(unsafe.Pointer(c)))
	video = (c.capabilities & _V4L2_CAP_VIDEO_CAPTURE) == 0
	stream = (c.capabilities & _V4L2_CAP_STREAMING) == 0
	return
}

func (w *Webcam) GetPixelFormats() (m map[uint32]string) {
	m = make(map[uint32]string)
	for i := uint32(0); ; i++ {
		vfd := fmtdesc{
			index: i,
			t:     _V4L2_BUF_TYPE_VIDEO_CAPTURE,
		}
		err := ioctl(w.fd, _VIDIOC_ENUM_FMT, uintptr(unsafe.Pointer(&vfd)))
		if err != nil {
			return
		}
		var s string
		for _, b := range vfd.description {
			s = s + string(b)
		}
		m[vfd.pixelformat] = s
	}
	return
}

type FrameSize struct {
	MinWidth   uint32
	MaxWidth   uint32
	StepWidth  uint32
	MinHeight  uint32
	MaxHeight  uint32
	StepHeight uint32
}

func (s FrameSize) String() string {
	if s.StepWidth == 0 && s.StepHeight == 0 {
		return fmt.Sprintf("%dx%d", s.MaxWidth, s.MaxHeight)
	} else {
		return fmt.Sprintf("[%d-%d;%d]x[%d-%d;%d]", s.MinWidth, s.MaxWidth, s.StepWidth, s.MinHeight, s.MaxHeight, s.StepHeight)
	}
}

type FrameSizes []FrameSize

func (fs FrameSizes) Len() int { return len(fs) }
func (fs FrameSizes) Less(i, j int) bool {
	ls := fs[i].MaxWidth * fs[i].MaxHeight
	rs := fs[j].MaxWidth * fs[j].MaxHeight
	return ls < rs
}
func (fs FrameSizes) Swap(i, j int) { fs[i], fs[j] = fs[j], fs[i] }

func (w *Webcam) GetFrameSizes(code uint32) (fs FrameSizes) {
	var s FrameSize
	for i := uint32(0); ; i++ {
		vfse := frmsizeenum{
			index:        i,
			pixel_format: code,
		}
		if err := ioctl(w.fd, _VIDIOC_ENUM_FRAMESIZES, uintptr(unsafe.Pointer(&vfse))); err != nil {
			break
		}
		switch vfse.t {
		case _V4L2_FRMSIZE_TYPE_DISCRETE:
			s.MinWidth = vfse.frmsize[0]
			s.MaxWidth = vfse.frmsize[0]
			s.StepWidth = 0
			s.MinHeight = vfse.frmsize[1]
			s.MaxHeight = vfse.frmsize[1]
			s.StepHeight = 0
		case _V4L2_FRMSIZE_TYPE_CONTINUOUS, _V4L2_FRMSIZE_TYPE_STEPWISE:
			s.MinWidth = vfse.frmsize[0]
			s.MaxWidth = vfse.frmsize[1]
			s.StepWidth = vfse.frmsize[2]
			s.MinHeight = vfse.frmsize[3]
			s.MaxHeight = vfse.frmsize[4]
			s.StepHeight = vfse.frmsize[5]
		}
		fs = append(fs, s)
	}
	sort.Sort(fs)
	return
}

func (cam *Webcam) SetImageFormat(w, h, formatcode uint32) (uint32, uint32, uint32, error) {
	f := format{
		t: _V4L2_BUF_TYPE_VIDEO_CAPTURE,
		pix: pix{
			width:       w,
			height:      h,
			pixelformat: formatcode,
			field:       _V4L2_FIELD_ANY,
		},
	}
	err := ioctl(cam.fd, _VIDIOC_S_FMT, uintptr(unsafe.Pointer(&f)))
	return f.pix.pixelformat, f.pix.width, f.pix.height, err
}

func (w *Webcam) mmapRequestBuffers(bufcount uint32) (int, error) {
	req := requestbuffers{
		count:  bufcount,
		t:      _V4L2_BUF_TYPE_VIDEO_CAPTURE,
		memory: _V4L2_MEMORY_MMAP,
	}
	err := ioctl(w.fd, _VIDIOC_REQBUFS, uintptr(unsafe.Pointer(&req)))
	return int(req.count), err
}

func (w *Webcam) mmapQueryBuffer(index uint32) ([]byte, error) {
	buf := buffer{
		t:      _V4L2_BUF_TYPE_VIDEO_CAPTURE,
		memory: _V4L2_MEMORY_MMAP,
		index:  index,
	}
	err := ioctl(w.fd, _VIDIOC_QUERYBUF, uintptr(unsafe.Pointer(&buf)))
	if err != nil {
		return nil, err
	}
	return syscall.Mmap(int(w.fd), buf.offset, int(buf.length), syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_SHARED)
}

func (w *Webcam) mmapDequeueBuffer() (uint32, uint32, error) {
	buf := buffer{
		t:      _V4L2_BUF_TYPE_VIDEO_CAPTURE,
		memory: _V4L2_MEMORY_MMAP,
	}
	err := ioctl(w.fd, _VIDIOC_DQBUF, uintptr(unsafe.Pointer(&buf)))
	return buf.index, buf.bytesused, err
}

func (w *Webcam) mmapEnqueueBuffer(index uint32) error {
	buf := buffer{
		t:      _V4L2_BUF_TYPE_VIDEO_CAPTURE,
		memory: _V4L2_MEMORY_MMAP,
		index:  index,
	}
	return ioctl(w.fd, _VIDIOC_QBUF, uintptr(unsafe.Pointer(&buf)))
}

func (w *Webcam) mmapReleaseBuffers() error {
	for _, buf := range w.buffers {
		if err := syscall.Munmap(buf); err != nil {
			return err
		}
	}
	return nil
}

func (w *Webcam) StartStreaming() error {
	n, err := w.mmapRequestBuffers(uint32(256))
	if err != nil {
		return err
	}

	w.buffers = make([][]byte, n)
	for i := range w.buffers {
		buf, err := w.mmapQueryBuffer(uint32(i))
		if err != nil {
			return err
		}
		w.buffers[i] = buf
		err = w.mmapEnqueueBuffer(uint32(i))
		if err != nil {
			return err
		}
	}

	t := _V4L2_BUF_TYPE_VIDEO_CAPTURE
	return ioctl(w.fd, _VIDIOC_STREAMON, uintptr(unsafe.Pointer(&t)))
}

type TimeoutErr string

func (t TimeoutErr) Error() string {
	return string(t)
}

func (w *Webcam) WaitForFrame() (err error) {
	fds := new(goselect.FDSet)
	for {
		fds.Zero()
		fds.Set(w.fd)
		err = goselect.Select(int(w.fd)+1, fds, nil, nil, w.timeout)
		if err == nil {
			return
		}
		e, ok := err.(syscall.Errno)
		if ok {
			if e.Timeout() {
				return TimeoutErr("timeout")
			}
			if e.Temporary() {
				continue
			}
		}
	}
}

func (w *Webcam) ReadFrame() ([]byte, error) {
	i, l, err := w.mmapDequeueBuffer()
	if err != nil {
		return nil, err
	}

	buf := make([]byte, l)
	copy(buf, w.buffers[int(i)])

	err = w.mmapEnqueueBuffer(i)
	if err != nil {
		return nil, err
	}
	return buf, nil
}

func (w *Webcam) ReadJPEG() (image.Image, error) {
	buf, err := w.ReadFrame()
	if err != nil {
		return nil, err
	}
	return jpeg.Decode(bytes.NewBuffer(buf))
}

func (w *Webcam) ReadPNG() (image.Image, error) {
	buf, err := w.ReadFrame()
	if err != nil {
		return nil, err
	}
	return png.Decode(bytes.NewBuffer(buf))
}

func (w *Webcam) Close() error {
	if err := w.mmapReleaseBuffers(); err != nil {
		return err
	}
	return w.file.Close()
}