upscayl-server/upscayl/params.go

package upscayl

import (
	"fmt"
	"image"
	"os"
	"path/filepath"
	"strings"

	"trle5.xyz/upscayl-server/configs"
)

type Params struct {
	// input image path [jpg, png, webp]
	Input string
	// output image path [jpg, png, webp],
	// webp will be slower than jpg/png
	Output string
	// output image format, can be [jpg, png, webp],
	// webp will be slower than jpg/png.
	// default png or follow Output ext
	Format string
	// model name to use
	Model string
	// compression of the output image, can be 0 to 100, default is 0
	Compress int
	// custom output scale, default set to 1
	// some model may not support some scale
	Scale int
	// output image width, not work if Scale is set
	Width int
	// not work if Width == 0
	Height int
	// can be [box, triangle, cubicbspline, catmullrom, mitchell, pointsample].
	// empty will auto decide, not work if Width == 0 or Scale is set.
	// run `upscayl-bin -r help` for more info.
	ResizeFilter string
}

func (p *Params) Validate() error {
	if p.Input == "" || p.Output == "" {
		return fmt.Errorf("input and output must be set")
	}

	file, err := os.Open(p.Input)
	if err != nil {
		return fmt.Errorf("failed to open input image: %w", err)
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return fmt.Errorf("failed to get input image info: %w", err)
	}

	if stat.Size() == 0 {
		return fmt.Errorf("input image is empty")
	}

	_, format, err := image.DecodeConfig(file)
	if err != nil {
		return fmt.Errorf("failed to decode input image config: %w", err)
	}

	switch format {
	case "jpg", "jpeg", "png", "webp":
		// 支持的格式，不需要做任何处理
		break
	default:
		return fmt.Errorf("unsupported input image format: %s", format)
	}

	// 根据 Format 和 Output 判断是否支持格式
	switch strings.ToLower(p.Format) {
	case "jpg", "jpeg", "png", "webp":
		// 支持的格式，不需要做任何处理
		break
	case "":
		// 没有指定格式，根据输出文件名后缀判断是否支持
		strs := strings.Split(p.Output, ".")
		switch strings.ToLower(strs[len(strs)-1]) {
		case "jpg", "jpeg", "png", "webp":
			// 从输出文件名后缀判断出是支持的格式，不需要做任何处理
		default:
			// 不支持的格式
			return fmt.Errorf("output format not supported")
		}
	default:
		// 不支持的格式
		return fmt.Errorf("format not supported")
	}

	if p.Model != "" {
		m, err := os.Stat(filepath.Join(configs.ModelsDir, p.Model+".bin"))
		if err != nil || m.Size() == 0 {
			return fmt.Errorf("model bin not found or empty: %w", err)
		}
		p, err := os.Stat(filepath.Join(configs.ModelsDir, p.Model+".param"))
		if err != nil || p.Size() == 0 {
			return fmt.Errorf("model param not found or empty: %w", err)
		}
	} else if configs.DefaultModel != "" {
		p.Model = configs.DefaultModel
	} else {
		return fmt.Errorf("model required: no default model")
	}

	if p.Compress < 0 || p.Compress > 100 {
		return fmt.Errorf("compress must be between 0 and 100")
	}

	if p.Scale < 0 {
		return fmt.Errorf("scale must be greater than 0")
	}

	if p.Width < 0 {
		return fmt.Errorf("width must be greater than 0")
	}

	if p.Height < 0 {
		return fmt.Errorf("height must be greater than 0")
	}

	if p.Width != 0 || p.Height != 0 || p.ResizeFilter != "" {
		switch p.ResizeFilter {
		case "box", "triangle", "cubicbspline", "catmullrom", "mitchell", "pointsample":
			break
		case "":
			p.ResizeFilter = "default"
		default:
			return fmt.Errorf("invalid resize filter: %s", p.ResizeFilter)
		}

		if p.Scale != 0 {
			return fmt.Errorf("width, height, and resize_filter have no effect when scale is set")
		}
	} else if p.Scale == 0 {
		p.Scale = 1
	}

	return nil
}

func (p Params) ToArgs() []string {
	args := []string{}

	if p.Input != "" {
		args = append(args, "-i", p.Input)
	}

	if p.Output != "" {
		args = append(args, "-o", p.Output)
	}

	if p.Format != "" {
		args = append(args, "-f", p.Format)
	}

	if p.Model != "" {
		args = append(args, "-n", p.Model)
	}

	if p.Compress != 0 {
		args = append(args, "-c", fmt.Sprint(p.Compress))
	}

	if p.Scale != 0 {
		args = append(args, "-s", fmt.Sprint(p.Scale))
	}

	if p.Width != 0 && p.Height != 0 {
		if p.ResizeFilter != "" {
			args = append(args, "-r", fmt.Sprintf("%dx%d:%s", p.Width, p.Height, p.ResizeFilter))
		} else {
			args = append(args, "-r", fmt.Sprintf("%d:%d", p.Width, p.Height))
		}
	} else if p.Width != 0 {
		if p.ResizeFilter != "" {
			args = append(args, "-w", fmt.Sprintf("%d:%s", p.Width, p.ResizeFilter))
		} else {
			args = append(args, "-w", fmt.Sprint(p.Width))
		}
	}

	return args
}

// SameParams checks if the values ​​of two Params parameters are equal (except for Input and Output).
// If any of the Params parameter Format is empty, the Output value suffix will be checked for comparison with Format; if this fails, false will be returned.
func SameParams(a, b Params) bool {
	if a.Format == "" || b.Format == "" {
		aStr := strings.Split(a.Output, ".")
		bStr := strings.Split(b.Output, ".")
		if len(aStr) == 0 || len(bStr) == 0 {
			return false
		}
		a.Format = strings.ToLower(aStr[len(aStr)-1])
		b.Format = strings.ToLower(bStr[len(bStr)-1])
	}

	return a.Format       == b.Format &&
		   a.Model        == b.Model &&
		   a.Compress     == b.Compress &&
		   a.Scale        == b.Scale &&
		   a.Width        == b.Width &&
		   a.Height       == b.Height &&
		   a.ResizeFilter == b.ResizeFilter
}