gbase/ocr/preprocess.go

package ocr

import (
	"fmt"
	"github.com/sirupsen/logrus"
	"image"
	"image/jpeg"
	"os"
	"path/filepath"
	"strings"
)

func Preprocess(rawPath string, stepList []string) (processedPath string, err error) {

	steps := strings.Join(stepList, ",")
	steps = strings.ToLower(steps)

	handle, err := os.Open(rawPath)
	if err != nil {
		logrus.Errorf("os.Open rawPath(%s) error : %v\n", rawPath, err)
		return processedPath, err
	}

	defer handle.Close()

	src, err := jpeg.Decode(handle)
	if err != nil {
		logrus.Errorf("jpeg.Encode error : %v\n", err)
		return processedPath, err
	}
	var grayData *image.Gray

	if strings.Contains(steps, "gray") == true {
		grayData = Gray(src, ToGrayLuminance)

		// Apply binarization using Otsu's method
		grayImg := image.NewGray(grayData.Bounds())
		for y := 0; y < grayData.Bounds().Dy(); y++ {
			for x := 0; x < grayData.Bounds().Dx(); x++ {
				grayImg.Set(x, y, grayData.At(x, y))
			}
		}
	} else if strings.Contains(steps, "threshold") == true {
		threshold := Otsu(grayData)
		Threshold(grayData, threshold, 255, 0)
	}

	rawDir, rawFileName := filepath.Split(rawPath)
	rawExt := filepath.Ext(rawFileName)
	processedPath = filepath.Join(rawDir, fmt.Sprintf("%s_processed%v", rawFileName[:len(rawFileName)-len(rawExt)], rawExt))

	processedHandle, err := os.Create(processedPath)
	if err != nil {
		logrus.Errorf("os.Open processedPath(%s) error : %v\n", processedPath, err)
		return processedPath, err
	}

	defer processedHandle.Close()

	err = jpeg.Encode(processedHandle, grayData, &jpeg.Options{Quality: 80})
	if err != nil {
		logrus.Errorf("jpeg.Encode error : %v\n", err)
		return processedPath, err
	}

	return processedPath, nil
}