fyp/logic/models/add.go

package models

import (
	"bytes"
	"image"
	"image/color"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"net/http"
	"os"
	"path"

	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/utils"
)

func loadBaseImage(c *Context, id string) {
	// TODO handle more types than png
	infile, err := os.Open(path.Join("savedData", id, "baseimage.png"))
	if err != nil {
		c.Logger.Errorf("Failed to read image for model with id %s\n", id)
		c.Logger.Error(err)
		ModelUpdateStatus(c, id, FAILED_PREPARING)
		return
	}
	defer infile.Close()

	src, format, err := image.Decode(infile)
	if err != nil {
		c.Logger.Errorf("Failed to decode image for model with id %s\n", id)
		c.Logger.Error(err)
		ModelUpdateStatus(c, id, FAILED_PREPARING)
		return
	}
	switch format {
	case "png":
	case "jpeg":
	default:
		ModelUpdateStatus(c, id, FAILED_PREPARING)
		c.Logger.Error("Found unkown format '%s'\n", "format", format)
		panic("Handle diferent files than .png")
	}

	var model_color string

	bounds := src.Bounds()
	width, height := bounds.Max.X, bounds.Max.Y

	switch src.ColorModel() {
	case color.Gray16Model:
		fallthrough
	case color.GrayModel:
		model_color = "greyscale"
	case color.NRGBAModel:
		fallthrough
	case color.RGBAModel:
		fallthrough
	case color.YCbCrModel:
		model_color = "rgb"
	default:
		c.Logger.Error("Do not know how to handle this color model")

		if src.ColorModel() == color.RGBA64Model {
			c.Logger.Error("Color is rgb 64")
		} else if src.ColorModel() == color.NRGBA64Model {
			c.Logger.Error("Color is nrgb 64")
		} else if src.ColorModel() == color.AlphaModel {
			c.Logger.Error("Color is alpha")
		} else if src.ColorModel() == color.CMYKModel {
			c.Logger.Error("Color is cmyk")
		} else if src.ColorModel() == color.NRGBA64Model {
			c.Logger.Error("Color is cmyk")
		} else if src.ColorModel() == color.NYCbCrAModel {
			c.Logger.Error("Color is cmyk a")
		} else if src.ColorModel() == color.Alpha16Model {
			c.Logger.Error("Color is cmyk a")
		} else {
			c.Logger.Error("Other so assuming color", "color mode", src.ColorModel())
		}

		ModelUpdateStatus(c, id, FAILED_PREPARING)
		return
	}

	// Note: this also updates the status to 2
	_, err = c.Db.Exec("update models set width=$1, height=$2, color_mode=$3, format=$4, status=$5 where id=$6", width, height, model_color, format, CONFIRM_PRE_TRAINING, id)
	if err != nil {
		c.Logger.Error("Could not update model")
		c.Logger.Error(err)
		ModelUpdateStatus(c, id, -1)
		return
	}
}

func handleAdd(handle *Handle) {
	handle.Post("/models/add", func(c *Context) *Error {
		if !c.CheckAuthLevel(1) {
			return nil
		}

		read_form, err := c.R.MultipartReader()
		if err != nil {
			return c.JsonErrorBadRequest(err, "Please provide a valid multipart Reader request")
		}

		var name string
		var file []byte

		for {
			part, err_part := read_form.NextPart()
			if err_part == io.EOF {
				break
			} else if err_part != nil {
				return &Error{Code: http.StatusBadRequest}
			}
			if part.FormName() == "name" {
				buf := new(bytes.Buffer)
				buf.ReadFrom(part)
				name = buf.String()
			}
			if part.FormName() == "file" {
				buf := new(bytes.Buffer)
				buf.ReadFrom(part)
				file = buf.Bytes()
			}
		}

		if name == "" || len(file) == 0 {
			return c.JsonBadRequest("Name is empty or file is empty")
		}

		row, err := handle.Db.Query("select id from models where name=$1 and user_id=$2;", name, c.User.Id)
		if err != nil {
			return c.Error500(err)
		}

		if row.Next() {
			return c.JsonBadRequest("Model with that name already exists!")
		}

		row, err = handle.Db.Query("insert into models (user_id, name) values ($1, $2) returning id", c.User.Id, name)
		if err != nil || !row.Next() {
			return c.Error500(err)
		}

		var id string
		err = row.Scan(&id)
		if err != nil {
			return c.Error500(err)
		}

		// TODO mk this path configurable
		dir_path := path.Join("savedData", id)

		err = os.Mkdir(dir_path, os.ModePerm)
		if err != nil {
			return c.Error500(err)
		}

		f, err := os.Create(path.Join(dir_path, "baseimage.png"))
		if err != nil {
			return c.Error500(err)
		}
		defer f.Close()

		f.Write(file)

		c.Logger.Warn("Created model with id %s! Started to proccess image!\n", "id", id)
		go loadBaseImage(c, id)

		return c.SendJSON(id)
	})
}