stash-bot/src/ihash/ihash.go
2023-07-05 10:37:32 +03:00

63 lines
1.7 KiB
Go

package ihash
import (
"encoding/binary"
"fmt"
"image"
"github.com/corona10/goimagehash"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/bsontype"
)
type IHash struct {
pHash, dHash, aHash *goimagehash.ImageHash
}
func (h IHash) Compare(other IHash) uint64 {
pd, _ := h.pHash.Distance(other.pHash)
dd, _ := h.dHash.Distance(other.dHash)
ad, _ := h.aHash.Distance(other.aHash)
return uint64(pd + dd + ad)
}
func GenerateIHash(img image.Image) (IHash, error) {
hash := IHash{}
var err error
hash.pHash, err = goimagehash.PerceptionHash(img)
if err != nil {
return hash, fmt.Errorf("perception hash error: %s", err.Error())
}
hash.dHash, err = goimagehash.DifferenceHash(img)
if err != nil {
return hash, fmt.Errorf("difference hash error: %s", err.Error())
}
hash.aHash, err = goimagehash.AverageHash(img)
if err != nil {
return hash, fmt.Errorf("average hash error: %s", err.Error())
}
return hash, nil
}
func (h IHash) MarshalBSONValue() (bsontype.Type, []byte, error) {
b := make([]byte, 0, 8*3)
b = binary.LittleEndian.AppendUint64(b, h.pHash.GetHash())
b = binary.LittleEndian.AppendUint64(b, h.dHash.GetHash())
b = binary.LittleEndian.AppendUint64(b, h.aHash.GetHash())
return bson.MarshalValue(b)
}
// UnmarshalBSONValue implements bson.ValueUnmarshaler interface.
func (h *IHash) UnmarshalBSONValue(t bsontype.Type, data []byte) error {
if h == nil {
return bson.ErrDecodeToNil
}
h.pHash = goimagehash.NewImageHash(binary.LittleEndian.Uint64(data[0:8]), goimagehash.PHash)
h.dHash = goimagehash.NewImageHash(binary.LittleEndian.Uint64(data[8:16]), goimagehash.DHash)
h.aHash = goimagehash.NewImageHash(binary.LittleEndian.Uint64(data[16:24]), goimagehash.PHash)
return nil
}