initial commit
This commit is contained in:
commit
dd48aef782
8 changed files with 1010 additions and 0 deletions
310
main.go
Normal file
310
main.go
Normal file
|
|
@ -0,0 +1,310 @@
|
|||
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())
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue