sauce-finder/main.go

package main

import (
	"bytes"
	"cmp"
	"context"
	"errors"
	"fmt"
	"image/color"
	// "sort"

	// "image/png"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"sauce/shared"
	"slices"
	"time"

	"gocv.io/x/gocv"
	"gocv.io/x/gocv/contrib"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

var (
	port    = ":9393"
	hashes  []contrib.ImgHashBase
	matcher = gocv.NewBFMatcher()
	phash   = contrib.PHash{}
	avghash = contrib.AverageHash{}
	// index = loadIndex()
	db *gorm.DB
)

type candidate struct {
	page            shared.Page
	phash, avghash  float64
	matches         []gocv.DMatch
	averageDistance float64
	publication     shared.Publication
}

func loadImageFromDisk(path string) (shared.Page, error) {
	img := gocv.IMRead(path, gocv.IMReadColor)
	if img.Empty() {
		log.Panic("cannot read image", path)
	}

	return shared.LoadImage(path, img)
}

// func (e *shared.Page) hashImage() {
//     phash.Compute(e.image, &e.phash)
//     if e.phash.Empty() {
//         panic("empty")
//     }
//     avghash.Compute(e.image, &e.avghash)
//     if e.phash.Empty() {
//         panic("empty")
//     }
// }

func newCandidate(e shared.Page, p shared.Publication) candidate {
	return candidate{
		page:        e,
		publication: p,
		// phash: phash.Compare(e.phash, b.phash) / 64,
		// avghash: avghash.Compare(e.avghash, b.avghash) / 64,
	}
}

func (c *candidate) tryMatch(search shared.Page) {
	c.matches = matcher.Match(search.Descriptors, c.page.Descriptors)

	slices.SortFunc(c.matches, func(a, b gocv.DMatch) int {
		return cmp.Compare(a.Distance, b.Distance)
	})

	var average float64
	for _, m := range c.matches {
		average += m.Distance
	}

	c.averageDistance = average / float64(len(c.matches))
}

// todo: paralelizar
func loadIndex() []shared.Publication {

	now := time.Now()
	log.Println("loading index...")

	const indexFolder = "index"

	indexDir, err := os.ReadDir(indexFolder)
	if err != nil {
		panic(err)
	}

	var index []shared.Publication

	for _, i := range indexDir {
		if !i.Type().IsDir() {
			continue
		}

		var pages []shared.Page

		cachePath := path.Join(indexFolder, i.Name(), "cache")
		pagesPath := path.Join(indexFolder, i.Name(), "pages")

		_, err := os.Stat(cachePath) // validade cache
		if errors.Is(err, os.ErrNotExist) {

			pagesFolder, err := os.ReadDir(pagesPath)
			if err != nil {
				log.Println(err)
				continue
			}

			err = os.Mkdir(cachePath, os.ModePerm)
			if err != nil {
				panic(err)
			}

			for _, p := range pagesFolder {

				e, err := loadImageFromDisk(path.Join(pagesPath, p.Name()))
				if err != nil {
					log.Println(err)
					continue
				}

				e.SaveORBtoDisk(path.Join(cachePath, p.Name()))

				pages = append(pages, e)

				// img, err := e.Descriptors.ToImage()
				// if err != nil {
				//     panic(err)
				// }
				//
				// cache, err := os.Create(path.Join(cachePath, p.Name()))
				// if err != nil {
				//     panic(err)
				// }
				// err = png.Encode(cache, img)

			}

		} else if err != nil {
			panic(err)

		} else {

			cacheDir, err := os.ReadDir(cachePath)
			if err != nil {
				panic(err)
			}

			for _, c := range cacheDir {
				des := gocv.IMRead(path.Join(cachePath, c.Name()), gocv.IMReadAnyColor)

				pages = append(pages, shared.Page{
					Descriptors: des,
					Path:        path.Join(pagesPath, c.Name()),
					Name:        c.Name(),
				})
			}
		}

		index = append(index, shared.Publication{
			Title: i.Name(),
			Pages: pages,
		})
	}

	log.Println("index loaded in", time.Since(now))
	return index
}

func drawMatches(a, b shared.Page, matches []gocv.DMatch, path string) {
	output := gocv.NewMat()
	gocv.DrawMatches(
		a.Image, a.Keypoints,
		b.Image, b.Keypoints,
		matches[:20],
		&output,
		color.RGBA{R: 255}, color.RGBA{R: 255}, nil,
		gocv.NotDrawSinglePoints,
	)
	gocv.IMWrite(path, output)
	// fmt.Println()
	// img2 := gocv.NewMat()
	// gocv.DrawKeyPoints(search.image, kp, &img2, color.RGBA{R: 255}, 0)
	// gocv.IMWrite("matches.png", img3)
}

func handleSearch(w http.ResponseWriter, req *http.Request) {

	fileReader, _, err := req.FormFile("search")
	if err != nil {
		panic(err)
	}
	file, err := io.ReadAll(fileReader)
	if err != nil {
		panic(err)
	}

	search, err := shared.LoadImageFromBytes(file)

	var candidates []candidate

	rows, err := db.Debug().Model(&shared.Page{}).Preload("publications").Rows()
	if err != nil {
		panic(err)
	}
	defer rows.Close()

	for rows.Next() {

		var page shared.Page
		db.ScanRows(rows, &page)

		page.Descriptors, err = gocv.NewMatFromBytes(500, 32, gocv.MatTypeCV8U, page.DescriptorBlob)
		if err != nil {
			panic(err)
		}

		c := newCandidate(page, shared.Publication{})

		c.tryMatch(search)

		candidates = append(candidates, c)
		// }
	}

	slices.SortFunc(candidates, func(a, b candidate) int {
		return cmp.Compare(a.averageDistance, b.averageDistance)
	})

	var pages []shared.Page

	for _, c := range candidates[:8] {
		var pub shared.Publication
		err = db.Where("id = ?", c.page.UserID).Find(&pub).Error
		if err != nil {
			panic(err)
		}
		fmt.Println("pub:", pub)
		c.page.Publication = pub
		pages = append(pages, c.page)
	}

	layout(results(search.B64, pages)).Render(context.Background(), w)
}

func main() {

	var err error
	db, err = gorm.Open(sqlite.Open("test.db"), &gorm.Config{})
	if err != nil {
		panic(err)
	}

	home := bytes.Buffer{}
	layout(form()).Render(context.Background(), &home)

	router := http.NewServeMux()

	router.HandleFunc("GET /", func(w http.ResponseWriter, req *http.Request) {
		w.Write(home.Bytes())
	})

	router.HandleFunc("POST /search", handleSearch)

	router.HandleFunc("GET /src", func(w http.ResponseWriter, r *http.Request) {
		url := r.FormValue("src")
		if url == "" {
			panic(url)
		}

		resp1, err := http.Get(url)
		if err != nil {
			panic(err)
		}
		if resp1.StatusCode != 200 {
			panic(resp1.Status)
		}
		defer resp1.Body.Close()

		_, err = io.Copy(w, resp1.Body)
		if err != nil {
			panic(err)
		}
	})

	router.Handle("GET /index/", http.StripPrefix("/index/", http.FileServer(http.Dir("index"))))

	server := http.Server{
		Addr:    port,
		Handler: Logging(router),
	}

	fmt.Println("http://localhost" + port)
	log.Fatal(server.ListenAndServe())
}