github.com/amitbet/vnc2video@v0.0.0-20190616012314-9d50b9dab1d9/encoders/mjpeg-enc.go (about)

     1  package encoders
     2  
     3  import (
     4  	"bytes"
     5  	"image"
     6  	"image/jpeg"
     7  	"strings"
     8  	"github.com/amitbet/vnc2video/logger"
     9  
    10  	"github.com/icza/mjpeg"
    11  )
    12  
    13  type MJPegImageEncoder struct {
    14  	avWriter  mjpeg.AviWriter
    15  	Quality   int
    16  	Framerate int32
    17  	closed    bool
    18  }
    19  
    20  func (enc *MJPegImageEncoder) Init(videoFileName string) {
    21  	fileExt := ".avi"
    22  	if enc.Framerate == 0 {
    23  		enc.Framerate = 12
    24  	}
    25  	if !strings.HasSuffix(videoFileName, fileExt) {
    26  		videoFileName = videoFileName + fileExt
    27  	}
    28  	if enc.Framerate <= 0 {
    29  		enc.Framerate = 5
    30  	}
    31  	avWriter, err := mjpeg.New(videoFileName, 1024, 768, enc.Framerate)
    32  	if err != nil {
    33  		logger.Error("Error during mjpeg init: ", err)
    34  	}
    35  	enc.avWriter = avWriter
    36  }
    37  func (enc *MJPegImageEncoder) Run(videoFileName string) {
    38  	enc.Init(videoFileName)
    39  }
    40  
    41  func (enc *MJPegImageEncoder) Encode(img image.Image) {
    42  	if enc.closed {
    43  		return
    44  	}
    45  
    46  	buf := &bytes.Buffer{}
    47  	jOpts := &jpeg.Options{Quality: enc.Quality}
    48  	if enc.Quality <= 0 {
    49  		jOpts = nil
    50  	}
    51  	err := jpeg.Encode(buf, img, jOpts)
    52  	if err != nil {
    53  		logger.Error("Error while creating jpeg: ", err)
    54  	}
    55  
    56  	//logger.Tracef("buff: %v\n", buf.Bytes())
    57  
    58  	err = enc.avWriter.AddFrame(buf.Bytes())
    59  	if err != nil {
    60  		logger.Error("Error while adding frame to mjpeg: ", err)
    61  	}
    62  }
    63  
    64  func (enc *MJPegImageEncoder) Close() {
    65  	err := enc.avWriter.Close()
    66  
    67  	enc.closed = true
    68  	if err != nil {
    69  		logger.Error("Error while closing mjpeg: ", err)
    70  	}
    71  }