-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.go
115 lines (90 loc) · 2.29 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
package main
import (
"fmt"
"image"
"image/color"
"image/jpeg"
"os"
"strconv"
"github.com/KimJeongChul/nvpipe"
"gocv.io/x/gocv"
)
func main() {
codec := nvpipe.NvPipeH264
if codec == nvpipe.NvPipeH264 {
fmt.Println("NvPipe_Codec : H.264")
} else {
fmt.Println("NvPipe_Codec : HEVC")
}
compression := nvpipe.NvPipeLossless
format := nvpipe.NvPipeRGBA32
width := 1280
height := 720
rgbaChannel := 4
fmt.Println("width : ", width, "height : ", height)
bitrateMbps := 32
targetFPS := 30
fmt.Println("bitrate : ", bitrateMbps, " targetFPS : ", targetFPS)
// NVDEC
decoder := nvpipe.NewDecoder(format, codec, width, height)
// NVENC
encoder := nvpipe.NewEncoder(format, codec, compression, bitrateMbps, targetFPS, width, height)
video, err := gocv.OpenVideoCapture("input.mp4") // h264 video file
if err != nil {
fmt.Println("[ERROR] gocv open video capture error : ", err)
}
defer video.Close()
mat := gocv.NewMat()
defer mat.Close()
fileIdx := 0
for {
if ok := video.Read(&mat); !ok {
return
}
if mat.Empty() {
break
}
img, err := mat.ToImage() // RGB 3 channel
rect := img.Bounds()
rgba := image.NewRGBA(rect)
// NvPipe Encoding
encodeSize := width * height
encodeData := make([]byte, encodeSize)
n := encoder.Encode(rgba.Pix, encodeData)
if n == 0 {
fmt.Println("[ERROR] nvenc error")
}
// Nvpipe Decoding
decodeData := make([]uint8, width*height*rgbaChannel) // 1280 * 720 * 4 = 3686400
n = decoder.Decode(encodeData[:n], n, decodeData)
if n == 0 {
fmt.Println("[ERROR] nvdec error")
}
jpgFile, err := os.Create("output" + strconv.Itoa(fileIdx) + ".jpg")
if err != nil {
fmt.Println("[ERROR] image file create : ", err)
}
// Preprocessing using image.RGBA
start := image.Point{0, 0}
end := image.Point{width, height}
output := image.NewRGBA(image.Rectangle{start, end})
// RGBA
rgbaSlice := make([]color.RGBA, 0)
for i := 0; i < width*height; i++ {
r := decodeData[4*i+0]
g := decodeData[4*i+1]
b := decodeData[4*i+2]
a := decodeData[4*i+3]
rgbaSlice = append(rgbaSlice, color.RGBA{r, g, b, a})
}
// Draw RGBA
for y := 0; y < height; y++ {
for x := 0; x < width; x++ {
output.Set(x, y, rgbaSlice[y*width+x])
}
}
// File Write
jpeg.Encode(jpgFile, output, nil)
fileIdx++
}
}