将图片数据转换为张量（Go并发处理）

拙_言2024-04-15 9:46

在Go语言中，将图片数据转换成Tensor通常需要依赖一些外部库，编写一个简单的程序，该程序批量同时处理图片，将其转换为对应的浮点数张量。

假设图片是单通道（灰度图）或者三通道（彩色图），我们将图片的每个像素值归一化到0到1之间，然后创建一个gorgonia张量。

注意sync.WaitGroup和sync.Mutex的使用

Go 复制代码

package main

import (
	"fmt"
	"image"
	"image/png"
	"os"
	"sync"

	"gorgonia.org/gorgonia"
)

func processImage(imagePath string, tensor *gorgonia.Tensor, wg *sync.WaitGroup, lock *sync.Mutex) error {
	img, err := os.Open(imagePath)
	if err != nil {
		return fmt.Errorf("error opening image file: %v", err)
	}
	defer img.Close()

	imgBounds := img.Bounds()
	img = image.NewRGBA(imgBounds)
	err = png.Decode(img, img)
	if err != nil {
		return fmt.Errorf("error decoding image: %v", err)
	}

	lock.Lock()
	defer lock.Unlock()

	for y := imgBounds.Min.Y; y < imgBounds.Max.Y; y++ {
		for x := imgBounds.Min.X; x < imgBounds.Max.X; x++ {
			r, g, b, _ := img.At(x, y).RGBA()
			tensor.SetWithShape(y, x, 0, float32(r)/255.0) // 设置红色通道
			tensor.SetWithShape(y, x, 1, float32(g)/255.0) // 设置绿色通道
			tensor.SetWithShape(y, x, 2, float32(b)/255.0) // 设置蓝色通道
		}
	}

	return nil // 没有错误发生
}

func main() {
	var wg sync.WaitGroup
	var lock sync.Mutex
	imagePaths := []string{
		"path/xcl/image1.png",
		"path/xcl/image2.png",
		"path/xcl/image3.png",
	}

	// 假设所有图片大小和通道数相同
	img, err := os.Open(imagePaths[0])
	if err != nil {
		panic(err)
	}
	defer img.Close()

	imgBounds := img.Bounds()
	tensorShape := gorgonia.WithShape(imgBounds.Dy(), imgBounds.Dx(), 3) // 假设RGB图片
	tensor := gorgonia.NewTensor(tensorShape)

	for _, imagePath := range imagePaths {
		wg.Add(1)
		go func(imagePath string, tensor *gorgonia.Tensor, wg *sync.WaitGroup, lock *sync.Mutex) {
			defer wg.Done()
			err := processImage(imagePath, tensor, wg, &lock)
			if err != nil {
				fmt.Printf("Error processing image %s: %v\n", imagePath, err)
			}
		}(imagePath, tensor, &wg, &lock)
	}

	// 等待所有goroutines完成
	wg.Wait()

	// 在这里检查是否有任何错误发生
	for _, imagePath := range imagePaths {
		fmt.Printf("Processed image: %s\n", imagePath)
	}
	fmt.Println(tensor)
}

请注意，这个例子假设图片是PNG格式的，并且是RGB彩色图片。

如果实际中，图片是其他格式或颜色模式，那需要另外再调整代码。