✨ feat: rewrite in Go

2026-09-13 18:59:29 +00:00 · 2026-07-13 17:58:48 +07:00
parent a7d5177100
commit 1421b854d4
34 changed files with 1714 additions and 1825 deletions
@@ -0,0 +1,179 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"math"
+	"os/exec"
+	"strconv"
+	"strings"
+)
+
+var (
+	jpegSOI = []byte{0xff, 0xd8}
+	jpegEOI = []byte{0xff, 0xd9}
+)
+
+type jpegFrameSplitter struct {
+	buffer []byte
+}
+
+func (s *jpegFrameSplitter) push(chunk []byte) [][]byte {
+	s.buffer = append(s.buffer, chunk...)
+	var frames [][]byte
+	for {
+		start := bytes.Index(s.buffer, jpegSOI)
+		if start == -1 {
+			break
+		}
+		end := bytes.Index(s.buffer[start+len(jpegSOI):], jpegEOI)
+		if end == -1 {
+			break
+		}
+		frameEnd := start + len(jpegSOI) + end + len(jpegEOI)
+		frame := make([]byte, frameEnd-start)
+		copy(frame, s.buffer[start:frameEnd])
+		frames = append(frames, frame)
+		s.buffer = s.buffer[frameEnd:]
+	}
+	return frames
+}
+
+type ffprobeOutput struct {
+	Streams []struct {
+		RFrameRate string `json:"r_frame_rate"`
+		NbFrames   string `json:"nb_frames"`
+		Duration   string `json:"duration"`
+	} `json:"streams"`
+	Format struct {
+		Duration string `json:"duration"`
+	} `json:"format"`
+}
+
+func parseFrameRate(rate string) float64 {
+	if rate == "" {
+		return math.NaN()
+	}
+	parts := strings.SplitN(rate, "/", 2)
+	num, err := strconv.ParseFloat(parts[0], 64)
+	if err != nil {
+		return math.NaN()
+	}
+	if len(parts) == 2 {
+		den, err := strconv.ParseFloat(parts[1], 64)
+		if err != nil || den == 0 {
+			return math.NaN()
+		}
+		return num / den
+	}
+	return num
+}
+
+func processVideo(path, output string, maxDimension int) error {
+	probeCmd := exec.Command(
+		"ffprobe", "-v", "error",
+		"-show_streams", "-show_format",
+		"-of", "json", path,
+	)
+	probeOut, err := probeCmd.Output()
+	if err != nil {
+		msg := err.Error()
+		var exitErr *exec.ExitError
+		if errors.As(err, &exitErr) && len(exitErr.Stderr) > 0 {
+			msg = strings.TrimSpace(string(exitErr.Stderr))
+		}
+		return fmt.Errorf("ffprobe failed: %s", msg)
+	}
+
+	var probe ffprobeOutput
+	if err := json.Unmarshal(probeOut, &probe); err != nil {
+		return fmt.Errorf("ffprobe failed: %w", err)
+	}
+	if len(probe.Streams) == 0 {
+		return fmt.Errorf("unable to determine a valid video fps")
+	}
+	stream := probe.Streams[0]
+
+	fps := parseFrameRate(stream.RFrameRate)
+	if math.IsNaN(fps) || fps <= 0 {
+		return fmt.Errorf("unable to determine a valid video fps")
+	}
+
+	totalFrames := 0
+	if n, err := strconv.Atoi(stream.NbFrames); err == nil {
+		totalFrames = n
+	} else {
+		durStr := probe.Format.Duration
+		if durStr == "" {
+			durStr = stream.Duration
+		}
+		if d, err := strconv.ParseFloat(durStr, 64); err == nil {
+			totalFrames = int(math.Round(d * fps))
+		}
+	}
+
+	vf := fmt.Sprintf(
+		"format=gray,scale=w=%d:h=%d:force_original_aspect_ratio=decrease",
+		maxDimension, maxDimension,
+	)
+	cmd := exec.Command(
+		"ffmpeg", "-i", path,
+		"-c:v", "mjpeg",
+		"-q:v", "3",
+		"-vf", vf,
+		"-f", "image2pipe",
+		"pipe:1",
+	)
+	var stderr bytes.Buffer
+	cmd.Stderr = &stderr
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return fmt.Errorf("ffmpeg failed: %w", err)
+	}
+	if err := cmd.Start(); err != nil {
+		return fmt.Errorf("ffmpeg failed: %w", err)
+	}
+
+	videoData := FramesContainer{FPS: fps}
+	splitter := &jpegFrameSplitter{}
+	reportProgress := func(count int) {
+		if totalFrames > 0 {
+			pct := min(100, int(math.Round(float64(count)/float64(totalFrames)*100)))
+			fmt.Printf("\rGenerating frames: %d/%d (%d%%)", count, totalFrames, pct)
+		} else {
+			fmt.Printf("\rGenerating frames: %d", count)
+		}
+	}
+
+	buf := make([]byte, 256*1024)
+	for {
+		n, err := stdout.Read(buf)
+		if n > 0 {
+			videoData.Frames = append(videoData.Frames, splitter.push(buf[:n])...)
+			reportProgress(len(videoData.Frames))
+		}
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			cmd.Wait()
+			return fmt.Errorf("ffmpeg stream error: %s", err.Error())
+		}
+	}
+	if err := cmd.Wait(); err != nil {
+		return fmt.Errorf("ffmpeg failed: %s", strings.TrimSpace(stderr.String()))
+	}
+	if len(videoData.Frames) == 0 {
+		return fmt.Errorf("no frames were decoded from the video")
+	}
+
+	fmt.Println()
+	if err := writeTSF(output, &videoData); err != nil {
+		return err
+	}
+	logInfo(fmt.Sprintf("Saved %d frames to %s", len(videoData.Frames), output))
+	return nil
+}