package goconv
|
|
/*
|
#cgo CFLAGS: -I./ -I./inc -I/usr/local/cuda/include
|
#cgo CXXFLAGS: -I./ -I./inc -I/usr/local/cuda/include -std=c++11
|
#cgo LDFLAGS: -L/usr/local/cuda/lib64 -lnppi -lcudart -ldl
|
#include <stdlib.h>
|
#include "conv.h"
|
*/
|
import "C"
|
import (
|
"unsafe"
|
|
"basic.com/valib/godraw.git"
|
"basic.com/valib/gogpu.git"
|
"github.com/disintegration/imaging"
|
)
|
|
const (
|
need = 200
|
reserved = 512
|
)
|
|
func gpuIndex(lastIndex int) int {
|
indices := gogpu.RankGPU()
|
if len(indices) == 0 {
|
return -1
|
}
|
|
for _, v := range indices {
|
if v != lastIndex {
|
if gogpu.SatisfyGPU(v, need, need/2) {
|
return v
|
}
|
}
|
}
|
|
if gogpu.SatisfyGPU(lastIndex, need, reserved) {
|
return lastIndex
|
}
|
return -1
|
}
|
|
type convertor struct {
|
width int
|
height int
|
rWidth int
|
rHeight int
|
conv C.convHandle
|
}
|
|
var convts []*convertor
|
|
func find(w, h, rw, rh int) *convertor {
|
for _, v := range convts {
|
if v.width == w && v.height == h && v.rWidth == rw && v.rHeight == rh {
|
return v
|
}
|
}
|
gpu := gpuIndex(0)
|
if gpu < 0 {
|
return nil
|
}
|
cw := C.conv_create(C.int(w), C.int(h), C.int(rw), C.int(rh), C.int(gpu))
|
if cw == nil {
|
return nil
|
}
|
c := &convertor{w, h, rw, rh, cw}
|
convts = append(convts, c)
|
return c
|
}
|
|
// YUV2BGR yuv->bgr
|
func YUV2BGR(yuv []byte, w, h int) []byte {
|
|
cw := find(w, h, -1, -1)
|
if cw == nil {
|
return yuv2bgr(yuv, w, h)
|
}
|
var bgr *C.uchar
|
var bgrLen C.int
|
ret := C.yuv2bgr(cw.conv, unsafe.Pointer(&yuv[0]), &bgr, &bgrLen)
|
if ret != 0 {
|
return nil
|
}
|
const maxLen = 0x7fffffff
|
goBGRLen := int(bgrLen)
|
if goBGRLen > 0 {
|
return (*[maxLen]byte)(unsafe.Pointer(bgr))[:goBGRLen:goBGRLen]
|
}
|
return nil
|
}
|
|
// YUV2ResizedBGR yuv -> resized bgr
|
func YUV2ResizedBGR(yuv []byte, w, h, rw, rh int) []byte {
|
|
cw := find(w, h, rw, rh)
|
if cw == nil {
|
bgr := yuv2bgr(yuv, w, h)
|
return bgresize(bgr, w, h, rw, rh)
|
}
|
var bgr *C.uchar
|
var bgrLen C.int
|
ret := C.yuv2resizedbgr(cw.conv, unsafe.Pointer(&yuv[0]), &bgr, &bgrLen)
|
if ret != 0 {
|
return nil
|
}
|
const maxLen = 0x7fffffff
|
goBGRLen := int(bgrLen)
|
if goBGRLen > 0 {
|
return (*[maxLen]byte)(unsafe.Pointer(bgr))[:goBGRLen:goBGRLen]
|
}
|
return nil
|
|
}
|
|
// ResizeBGR resize
|
func ResizeBGR(bgrO []byte, w, h, rw, rh int) []byte {
|
if (rw < 0 && rh < 0) || (rw > w && rh > h) {
|
return bgrO
|
}
|
|
cw := find(w, h, rw, rh)
|
if cw == nil {
|
return bgresize(bgrO, w, h, rw, rh)
|
}
|
|
var bgr *C.uchar
|
var bgrLen C.int
|
ret := C.resizebgr(cw.conv, unsafe.Pointer(&bgrO[0]), &bgr, &bgrLen)
|
if ret != 0 {
|
return nil
|
}
|
const maxLen = 0x7fffffff
|
goBGRLen := int(bgrLen)
|
if goBGRLen > 0 {
|
return (*[maxLen]byte)(unsafe.Pointer(bgr))[:goBGRLen:goBGRLen]
|
}
|
return nil
|
}
|
|
// ResizeYUV yuv
|
func ResizeYUV(yuv []byte, w, h, rw, rh int) []byte {
|
if (rw < 0 && rh < 0) || (rw > w && rh > h) {
|
return yuv
|
}
|
|
cw := find(w, h, rw, rh)
|
if cw == nil {
|
return yuv
|
}
|
|
var resized *C.uchar
|
var resizedLen C.int
|
ret := C.resizeyuv(cw.conv, unsafe.Pointer(&yuv[0]), &resized, &resizedLen)
|
if ret != 0 {
|
return nil
|
}
|
|
const maxLen = 0x7fffffff
|
goResizedLen := int(resizedLen)
|
if goResizedLen > 0 {
|
return (*[maxLen]byte)(unsafe.Pointer(resized))[:goResizedLen:goResizedLen]
|
}
|
return nil
|
}
|
|
// YUV2BGRandResize conv and resize
|
func YUV2BGRandResize(yuv []byte, w, h, rw, rh int) ([]byte, []byte) {
|
cw := find(w, h, rw, rh)
|
if cw == nil {
|
origin := yuv2bgr(yuv, w, h)
|
resized := bgresize(origin, w, h, rw, rh)
|
return origin, resized
|
}
|
|
var bgr *C.uchar
|
var bgrLen C.int
|
var scale *C.uchar
|
var scaleLen C.int
|
|
ret := C.yuv2bgrandresize(cw.conv, unsafe.Pointer(&yuv[0]), &bgr, &bgrLen, &scale, &scaleLen)
|
|
if ret != 0 {
|
return nil, nil
|
}
|
var out, resized []byte
|
|
const maxLen = 0x7fffffff
|
goBGRLen, goScaleLen := int(bgrLen), int(scaleLen)
|
if goBGRLen > 0 {
|
out = (*[maxLen]byte)(unsafe.Pointer(bgr))[:goBGRLen:goBGRLen]
|
}
|
if goScaleLen > 0 {
|
resized = (*[maxLen]byte)(unsafe.Pointer(scale))[:goScaleLen:goScaleLen]
|
}
|
return out, resized
|
|
}
|
|
// Free free
|
func Free() {
|
for _, v := range convts {
|
if v.conv != nil {
|
C.conv_destroy(v.conv)
|
}
|
}
|
}
|
|
func yuv2bgr(yuv []byte, w, h int) []byte {
|
|
data := make([]byte, 0, w*h*3)
|
start := w * h
|
for i := 0; i < h; i++ {
|
for j := 0; j < w; j++ {
|
|
index := i/2*w + j - (j & 0x01)
|
|
y := int32(yuv[j+i*w])
|
v := int32(yuv[start+index])
|
u := int32(yuv[start+index+1])
|
|
r := y + (140*(v-128))/100
|
g := y - (34*(u-128)+71*(v-128))/100
|
b := y + (177*(u-128))/100
|
|
if r > 255 {
|
r = 255
|
}
|
if r < 0 {
|
r = 0
|
}
|
if g > 255 {
|
g = 255
|
}
|
if g < 0 {
|
g = 0
|
}
|
if b > 255 {
|
b = 255
|
}
|
if b < 0 {
|
b = 0
|
}
|
data = append(data, byte(r), byte(g), byte(b))
|
}
|
}
|
return data
|
}
|
|
func bgresize(bgr []byte, w, h, rw, rh int) []byte {
|
img, err := godraw.ToImage(bgr, w, h)
|
if err != nil {
|
return nil
|
}
|
dstImg := imaging.Resize(img, rw, rh, imaging.NearestNeighbor)
|
return godraw.Image2BGR(dstImg)
|
}
|
