comic-hasher/hashing.go

160 lines
2.8 KiB
Go
Raw Normal View History

2024-08-04 18:12:00 -07:00
package ch
2024-05-01 18:09:02 -07:00
import (
2024-08-04 18:12:00 -07:00
"cmp"
_ "embed"
2024-05-01 18:09:02 -07:00
"fmt"
"image"
"log"
2024-08-04 18:12:00 -07:00
"math/bits"
2024-07-31 11:35:17 -07:00
"runtime"
2024-08-04 18:12:00 -07:00
"slices"
2024-05-01 18:09:02 -07:00
2024-07-31 11:35:17 -07:00
"gitea.narnian.us/lordwelch/goimagehash"
2024-05-01 18:09:02 -07:00
)
//go:embed hashes.gz
var Hashes []byte
2024-05-01 18:09:02 -07:00
const (
2024-08-04 18:26:20 -07:00
H0 uint64 = 0b11111111 << (8 * iota)
H1
H2
H3
H4
H5
H6
H7
2024-05-01 18:09:02 -07:00
)
const (
2024-08-04 18:26:20 -07:00
Shift0 = (8 * iota)
Shift1
Shift2
Shift3
Shift4
Shift5
Shift6
Shift7
2024-05-01 18:09:02 -07:00
)
2024-08-05 13:54:00 -07:00
const (
ComicVine Source = "comicvine.gamespot.com"
)
2024-08-04 18:12:00 -07:00
type Source string
2024-05-01 18:09:02 -07:00
2024-08-04 18:12:00 -07:00
type Match struct {
Distance int
Hash uint64
2024-05-01 18:09:02 -07:00
}
type ID struct {
Domain, ID string
}
2024-08-04 18:12:00 -07:00
type Result struct {
IDs []string // domain:id
2024-08-04 18:12:00 -07:00
Distance int
Hash ImageHash
2024-07-31 11:35:17 -07:00
}
2024-08-04 18:12:00 -07:00
type Im struct {
Im image.Image
Format string
Domain Source
ID, Path string
2024-07-31 11:35:17 -07:00
}
2024-08-04 18:12:00 -07:00
type Hash struct {
Ahash *goimagehash.ImageHash
Dhash *goimagehash.ImageHash
Phash *goimagehash.ImageHash
Domain Source
ID string
2024-05-01 18:09:02 -07:00
}
2024-08-04 18:12:00 -07:00
type ImageHash struct {
Hash uint64
Kind goimagehash.Kind
2024-05-01 18:09:02 -07:00
}
2024-08-04 18:26:20 -07:00
func Atleast(maxDistance int, searchHash uint64, hashes []uint64) []Match {
matchingHashes := make([]Match, 0, len(hashes)/2) // hope that we don't need all of them
for _, storedHash := range hashes {
distance := bits.OnesCount64(searchHash ^ storedHash)
2024-08-04 18:12:00 -07:00
if distance <= maxDistance {
2024-08-04 18:26:20 -07:00
matchingHashes = append(matchingHashes, Match{distance, storedHash})
2024-05-01 18:09:02 -07:00
}
}
2024-08-04 18:26:20 -07:00
return matchingHashes
2024-07-31 11:35:17 -07:00
}
2024-08-04 18:12:00 -07:00
func Insert[S ~[]E, E cmp.Ordered](slice S, item E) S {
2024-08-04 18:26:20 -07:00
index, itemFound := slices.BinarySearch(slice, item)
if itemFound {
2024-08-04 18:12:00 -07:00
return slice
2024-07-31 11:35:17 -07:00
}
2024-08-04 18:12:00 -07:00
return slices.Insert(slice, index, item)
2024-07-31 11:35:17 -07:00
}
func MemStats() uint64 {
var m runtime.MemStats
runtime.ReadMemStats(&m)
return m.Alloc
}
2024-08-04 18:12:00 -07:00
func HashImage(i Im) Hash {
if i.Format == "webp" {
i.Im = goimagehash.FancyUpscale(i.Im.(*image.YCbCr))
2024-07-31 11:35:17 -07:00
}
2024-05-01 18:09:02 -07:00
var (
2024-07-31 11:35:17 -07:00
err error = nil
2024-05-01 18:09:02 -07:00
ahash *goimagehash.ImageHash
dhash *goimagehash.ImageHash
phash *goimagehash.ImageHash
)
2024-08-04 18:12:00 -07:00
ahash, err = goimagehash.AverageHash(i.Im)
2024-05-01 18:09:02 -07:00
if err != nil {
msg := fmt.Sprintf("Failed to ahash Image: %s", err)
log.Println(msg)
2024-08-04 18:12:00 -07:00
return Hash{}
2024-05-01 18:09:02 -07:00
}
2024-08-04 18:12:00 -07:00
dhash, err = goimagehash.DifferenceHash(i.Im)
2024-05-01 18:09:02 -07:00
if err != nil {
msg := fmt.Sprintf("Failed to dhash Image: %s", err)
log.Println(msg)
2024-08-04 18:12:00 -07:00
return Hash{}
2024-05-01 18:09:02 -07:00
}
2024-08-04 18:12:00 -07:00
phash, err = goimagehash.PerceptionHash(i.Im)
2024-05-01 18:09:02 -07:00
if err != nil {
msg := fmt.Sprintf("Failed to phash Image: %s", err)
log.Println(msg)
2024-08-04 18:12:00 -07:00
return Hash{}
2024-05-01 18:09:02 -07:00
}
2024-08-04 18:12:00 -07:00
return Hash{
Ahash: ahash,
Dhash: dhash,
Phash: phash,
Domain: i.Domain,
ID: i.ID,
2024-05-01 18:09:02 -07:00
}
}
func SplitHash(hash uint64) [8]uint8 {
return [8]uint8{
2024-08-04 18:26:20 -07:00
uint8((hash & H7) >> Shift7),
uint8((hash & H6) >> Shift6),
uint8((hash & H5) >> Shift5),
uint8((hash & H4) >> Shift4),
uint8((hash & H3) >> Shift3),
uint8((hash & H2) >> Shift2),
uint8((hash & H1) >> Shift1),
uint8((hash & H0) >> Shift0),
2024-05-01 18:09:02 -07:00
}
}
2024-08-04 18:12:00 -07:00
type IDList map[Source][]string // IDs is a map of domain to ID eg IDs['comicvine.gamespot.com'] = []string{"1235"}