create the program
This commit is contained in:
315
cmd/cli/cli.go
Normal file
315
cmd/cli/cli.go
Normal file
@ -0,0 +1,315 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"image/color"
|
||||
"image/png"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/disintegration/imaging"
|
||||
"github.com/urfave/cli/v2"
|
||||
"golang.org/x/image/tiff"
|
||||
"image"
|
||||
"image/jpeg"
|
||||
"io"
|
||||
"mime"
|
||||
"os"
|
||||
"path"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func main() {
|
||||
commands := []*cli.Command{
|
||||
// add commands here
|
||||
}
|
||||
|
||||
app := &cli.App{
|
||||
Name: "imgStack",
|
||||
Usage: "A command line tool to denoise or give super resolution to an image by stacking multiple images together.",
|
||||
Version: "v1.0.0",
|
||||
Description: "imgStack is a command line tool to denoise or give super resolution to an image by stacking multiple images together.",
|
||||
Commands: commands,
|
||||
Flags: []cli.Flag{
|
||||
&cli.StringFlag{
|
||||
Name: "output",
|
||||
Value: "./output.png",
|
||||
Usage: "path to output file (should be a png)",
|
||||
Aliases: []string{
|
||||
"o",
|
||||
},
|
||||
},
|
||||
&cli.BoolFlag{
|
||||
Name: "super-resolution",
|
||||
Aliases: []string{"s", "sr"},
|
||||
Usage: "Whether to use super resolution or not. If this is set to true, the image will be doubled in size.",
|
||||
Value: false,
|
||||
},
|
||||
},
|
||||
Authors: []*cli.Author{
|
||||
{
|
||||
Name: "Mason Payne",
|
||||
Email: "mason@masonitestudios.com",
|
||||
},
|
||||
},
|
||||
Copyright: "2024 Masonite Studios LLC",
|
||||
UseShortOptionHandling: true,
|
||||
Action: func(c *cli.Context) error {
|
||||
var needsCleanup bool
|
||||
var cleanUpFiles []string
|
||||
var sourceFiles []string
|
||||
// use the first argument as the file name
|
||||
// we need to collect the list of discrete images
|
||||
fmt.Println("Args: ", c.Args().Slice())
|
||||
|
||||
fileList, err := getListOfFiles(c.Args().Slice())
|
||||
if err != nil {
|
||||
return fmt.Errorf("error getting list of files | %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("Files to be used: ", fileList)
|
||||
|
||||
// check if superResolution has been requested
|
||||
superResolution := c.Bool("super-resolution")
|
||||
|
||||
if len(fileList) == 0 {
|
||||
return fmt.Errorf("no files provided")
|
||||
}
|
||||
|
||||
// loop over each file
|
||||
for _, fileLocation := range fileList {
|
||||
// make sure the image is either jpeg, png or tiff
|
||||
fileMimeType := mime.TypeByExtension(path.Ext(fileLocation))
|
||||
if fileMimeType != "image/jpeg" && fileMimeType != "image/png" && fileMimeType != "image/tiff" {
|
||||
return fmt.Errorf("unsupported file type | %v", fileMimeType)
|
||||
}
|
||||
|
||||
if superResolution {
|
||||
// resize the images
|
||||
tempLocation, err := resizeImage(fileLocation, 2)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error resizing image | %w", err)
|
||||
}
|
||||
|
||||
if tempLocation != fileLocation {
|
||||
fmt.Printf("Resized image to %s\n", tempLocation)
|
||||
needsCleanup = true
|
||||
fileLocation = tempLocation
|
||||
}
|
||||
sourceFiles = append(sourceFiles, fileLocation) // use the resized image
|
||||
|
||||
fmt.Println("Needs cleanup: ", needsCleanup)
|
||||
if needsCleanup {
|
||||
cleanUpFiles = append(cleanUpFiles, tempLocation)
|
||||
}
|
||||
} else {
|
||||
sourceFiles = append(sourceFiles, fileLocation) // use the original image
|
||||
}
|
||||
}
|
||||
|
||||
err = stackImages(sourceFiles, c.String("output"))
|
||||
if err != nil {
|
||||
return fmt.Errorf("error stacking images | %w", err)
|
||||
}
|
||||
|
||||
if needsCleanup {
|
||||
for _, file := range cleanUpFiles {
|
||||
// remove the temp file
|
||||
err = cleanUpTempFile(file)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error removing temp file | %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
err := app.Run(os.Args)
|
||||
if err != nil {
|
||||
fmt.Println(fmt.Errorf("error running app | %w", err))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func cleanUpTempFile(fileLocation string) error {
|
||||
// remove the temp file
|
||||
err := os.Remove(fileLocation)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error removing temp file | %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func resizeImage(fileLocation string, multiple int) (string, error) {
|
||||
// load the image
|
||||
var img image.Image
|
||||
img, err := imaging.Open(fileLocation, imaging.AutoOrientation(true))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("error opening image | %w", err)
|
||||
}
|
||||
|
||||
// double the size of the image
|
||||
rect := img.Bounds()
|
||||
width := rect.Max.X - rect.Min.X
|
||||
height := rect.Max.Y - rect.Min.Y
|
||||
x1 := multiple * width
|
||||
y1 := multiple * height
|
||||
|
||||
croppedImage := imaging.Fill(img, x1, y1, imaging.Center, imaging.Linear) // technically not cropped because it will still have the same aspect ratio
|
||||
|
||||
// save the image to a temp file
|
||||
tempFile, err := os.CreateTemp("", "imgStack*"+path.Ext(fileLocation))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("error creating temp file | %w", err)
|
||||
}
|
||||
defer tempFile.Close()
|
||||
|
||||
// encode the image to the temp file
|
||||
fileMimeType := mime.TypeByExtension(path.Ext(fileLocation))
|
||||
if fileMimeType == "image/jpeg" {
|
||||
err = jpeg.Encode(tempFile, croppedImage, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("error encoding resized jpeg image | %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if fileMimeType == "image/png" {
|
||||
err = png.Encode(tempFile, croppedImage)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("error encoding resized png image | %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if fileMimeType == "image/tiff" {
|
||||
err = tiff.Encode(tempFile, croppedImage, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("error encoding resized tiff image | %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// return the temp file location
|
||||
return tempFile.Name(), nil
|
||||
}
|
||||
|
||||
// stackImages stacks the images using the average value of each pixel in sourceFiles and saves the result to output
|
||||
func stackImages(sourceFiles []string, output string) error {
|
||||
// load the images
|
||||
var images []image.Image
|
||||
for _, fileLocation := range sourceFiles {
|
||||
// load the image
|
||||
var img image.Image
|
||||
img, err := imaging.Open(fileLocation, imaging.AutoOrientation(true))
|
||||
if err != nil {
|
||||
return fmt.Errorf("error opening image | %w", err)
|
||||
}
|
||||
images = append(images, img)
|
||||
}
|
||||
|
||||
// get the bounds of the first image
|
||||
bounds := images[0].Bounds()
|
||||
|
||||
// create a new image
|
||||
newImage := image.NewRGBA(bounds)
|
||||
|
||||
// loop over each pixel in the new image
|
||||
for x := 0; x < bounds.Max.X; x++ {
|
||||
for y := 0; y < bounds.Max.Y; y++ {
|
||||
// get the average value of each pixel
|
||||
var r, g, b, a uint32
|
||||
for _, img := range images {
|
||||
// get the pixel at x, y
|
||||
pixel := img.At(x, y)
|
||||
// convert the pixel to RGBA
|
||||
_r, _g, _b, _a := pixel.RGBA()
|
||||
// get the values of the pixel
|
||||
r += _r
|
||||
g += _g
|
||||
b += _b
|
||||
a += _a
|
||||
}
|
||||
// get the average value of each pixel
|
||||
|
||||
r = r / uint32(len(images))
|
||||
g = g / uint32(len(images))
|
||||
b = b / uint32(len(images))
|
||||
a = a / uint32(len(images))
|
||||
// set the pixel at x, y
|
||||
|
||||
newImage.Set(x, y, &color.RGBA{
|
||||
R: uint8(r >> 8),
|
||||
G: uint8(g >> 8),
|
||||
B: uint8(b >> 8),
|
||||
A: uint8(a >> 8),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// save the new image
|
||||
err := saveImage(newImage, output)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error saving image | %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func saveImage(img image.Image, output string) error {
|
||||
// open the output file
|
||||
outputFile, err := os.Create(output)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error creating output file | %w", err)
|
||||
}
|
||||
defer outputFile.Close()
|
||||
|
||||
// encode the image to the output file
|
||||
fileMimeType := mime.TypeByExtension(path.Ext(output))
|
||||
if fileMimeType == "image/jpeg" {
|
||||
err = jpeg.Encode(outputFile, img, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error encoding resized jpeg image | %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if fileMimeType == "image/png" {
|
||||
err = png.Encode(outputFile, img)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error encoding resized png image | %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if fileMimeType == "image/tiff" {
|
||||
err = tiff.Encode(outputFile, img, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error encoding resized tiff image | %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func getListOfFiles(args []string) ([]string, error) {
|
||||
var filenames []string
|
||||
patterns := args
|
||||
|
||||
if len(patterns) == 0 {
|
||||
fmt.Println("No files provided, using stdin.")
|
||||
// read from stdin
|
||||
var err error
|
||||
fileLocationBytes, err := io.ReadAll(os.Stdin)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error reading from stdin | %w", err)
|
||||
}
|
||||
patterns = strings.Split(string(fileLocationBytes), "\n")
|
||||
}
|
||||
|
||||
for _, pattern := range patterns {
|
||||
subFileNames, err := filepath.Glob(pattern)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting list of files | %w", err)
|
||||
}
|
||||
filenames = append(filenames, subFileNames...)
|
||||
}
|
||||
|
||||
return filenames, nil
|
||||
}
|
||||
Reference in New Issue
Block a user