diff --git a/encoder.go b/encoder.go new file mode 100644 index 0000000..75f4470 --- /dev/null +++ b/encoder.go @@ -0,0 +1,180 @@ +package main + +/** +* XXX: Ugly workaround for https://github.com/amitbet/vnc2video/issues/10. I've copied the file and build a +* X264ImageCustomEncoder. Once this is merged, we can drop the encoder.go file again. +*/ + +import ( + "errors" + "fmt" + vnc "github.com/amitbet/vnc2video" + "github.com/amitbet/vnc2video/encoders" + log "github.com/sirupsen/logrus" + "image" + "image/color" + "io" + "os" + "os/exec" + "strconv" +) + +func encodePPMforRGBA(w io.Writer, img *image.RGBA) error { + maxvalue := 255 + size := img.Bounds() + // write ppm header + _, err := fmt.Fprintf(w, "P6\n%d %d\n%d\n", size.Dx(), size.Dy(), maxvalue) + if err != nil { + return err + } + + if convImage == nil { + convImage = make([]uint8, size.Dy()*size.Dx()*3) + } + + rowCount := 0 + for i := 0; i < len(img.Pix); i++ { + if (i % 4) != 3 { + convImage[rowCount] = img.Pix[i] + rowCount++ + } + } + + if _, err := w.Write(convImage); err != nil { + return err + } + + return nil +} + +func encodePPMGeneric(w io.Writer, img image.Image) error { + maxvalue := 255 + size := img.Bounds() + // write ppm header + _, err := fmt.Fprintf(w, "P6\n%d %d\n%d\n", size.Dx(), size.Dy(), maxvalue) + if err != nil { + return err + } + + // write the bitmap + colModel := color.RGBAModel + row := make([]uint8, size.Dx()*3) + for y := size.Min.Y; y < size.Max.Y; y++ { + i := 0 + for x := size.Min.X; x < size.Max.X; x++ { + color := colModel.Convert(img.At(x, y)).(color.RGBA) + row[i] = color.R + row[i+1] = color.G + row[i+2] = color.B + i += 3 + } + if _, err := w.Write(row); err != nil { + return err + } + } + return nil +} + +var convImage []uint8 + +func encodePPM(w io.Writer, img image.Image) error { + if img == nil { + return errors.New("nil image") + } + img1, isRGBImage := img.(*vnc.RGBImage) + img2, isRGBA := img.(*image.RGBA) + if isRGBImage { + return encodePPMforRGBImage(w, img1) + } else if isRGBA { + return encodePPMforRGBA(w, img2) + } + return encodePPMGeneric(w, img) +} +func encodePPMforRGBImage(w io.Writer, img *vnc.RGBImage) error { + maxvalue := 255 + size := img.Bounds() + // write ppm header + _, err := fmt.Fprintf(w, "P6\n%d %d\n%d\n", size.Dx(), size.Dy(), maxvalue) + if err != nil { + return err + } + + if _, err := w.Write(img.Pix); err != nil { + return err + } + return nil +} + +type X264ImageCustomEncoder struct { + encoders.X264ImageEncoder + FFMpegBinPath string + cmd *exec.Cmd + input io.WriteCloser + closed bool + Framerate int +} + +func (enc *X264ImageCustomEncoder) Init(videoFileName string) { + if enc.Framerate == 0 { + enc.Framerate = 12 + } + cmd := exec.Command(enc.FFMpegBinPath, + "-f", "image2pipe", + "-vcodec", "ppm", + "-r", strconv.Itoa(enc.Framerate), + "-an", //no audio + "-y", + "-i", "-", + "-vcodec", "libx264", + "-threads", "8", + "-preset", "veryfast", + "-g", "250", + "-crf", "37", + videoFileName, + ) + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + encInput, err := cmd.StdinPipe() + enc.input = encInput + if err != nil { + log.Error("can't get ffmpeg input pipe") + } + enc.cmd = cmd +} +func (enc *X264ImageCustomEncoder) Run(videoFileName string) error { + if _, err := os.Stat(enc.FFMpegBinPath); os.IsNotExist(err) { + if _, err := os.Stat(enc.FFMpegBinPath + ".exe"); os.IsNotExist(err) { + log.Error("encoder file doesn't exist in path:", enc.FFMpegBinPath) + return errors.New("encoder file doesn't exist in path" + videoFileName) + } else { + enc.FFMpegBinPath = enc.FFMpegBinPath + ".exe" + } + } + + enc.Init(videoFileName) + log.Infof("launching binary: %v", enc.cmd) + err := enc.cmd.Run() + if err != nil { + log.Errorf("error while launching ffmpeg: %v\n err: %v", enc.cmd.Args, err) + return err + } + return nil +} +func (enc *X264ImageCustomEncoder) Encode(img image.Image) { + if enc.input == nil || enc.closed { + return + } + + err := encodePPM(enc.input, img) + if err != nil { + log.Error("error while encoding image:", err) + } +} + +func (enc *X264ImageCustomEncoder) Close() { + if enc.closed { + return + } + enc.closed = true + enc.input.Close() +} diff --git a/main.go b/main.go index 43b7951..726d9eb 100644 --- a/main.go +++ b/main.go @@ -4,7 +4,6 @@ import ( "context" "fmt" vnc "github.com/amitbet/vnc2video" - "github.com/amitbet/vnc2video/encoders" log "github.com/sirupsen/logrus" cli "github.com/urfave/cli/v2" "net" @@ -141,7 +140,7 @@ func recorder(c *cli.Context) error { panic(err) } log.Infof("Using %s for encoding", ffmpeg_path) - vcodec := &encoders.X264ImageEncoder{ + vcodec := &X264ImageCustomEncoder{ FFMpegBinPath: ffmpeg_path, Framerate: c.Int("framerate"), } @@ -182,12 +181,16 @@ func recorder(c *cli.Context) error { } }() - sigc := make(chan os.Signal, 1) - signal.Notify(sigc, + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, + os.Interrupt, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, - syscall.SIGQUIT) + syscall.SIGQUIT, + syscall.SIGKILL, + ) + frameBufferReq := 0 timeStart := time.Now() @@ -202,22 +205,18 @@ func recorder(c *cli.Context) error { secsPassed := time.Now().Sub(timeStart).Seconds() frameBufferReq++ reqPerSec := float64(frameBufferReq) / secsPassed - //counter++ - //jpeg.Encode(out, screenImage, nil) - ///vcodec.Encode(screenImage) log.Debugf("reqs=%d, seconds=%f, Req Per second= %f", frameBufferReq, secsPassed, reqPerSec) reqMsg := vnc.FramebufferUpdateRequest{Inc: 1, X: 0, Y: 0, Width: cc.Width(), Height: cc.Height()} - //cc.ResetAllEncodings() reqMsg.Write(cc) } - case signal := <-sigc: + case signal := <-sigCh: if signal != nil { log.Info(signal, " received, exit.") vcodec.Close() // give some time to write the file - time.Sleep(time.Second * 5) - return nil + time.Sleep(time.Second * 1) + os.Exit(0) } } }