Mirrors/stash
1
0
Fork 0
mirror of https://github.com/stashapp/stash.git synced 2025-12-18 04:44:37 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add PHash distance matching to stash-box integration (#1858)

Browse Source 
* Add PHash distance matching to stash-box integration
This commit is contained in:
InfiniteTF
2021-10-20 08:22:25 +02:00
committed by GitHub
parent 976038424b
commit 15acf91b90
21 changed files with 462 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
pkg/scraper/stashbox/stash_box.go
View File

@@ -75,7 +75,7 @@ func (c Client) FindStashBoxScenesByFingerprints(sceneIDs []string) ([][]*models
return nil, err
}
var fingerprints []string
var fingerprints []*graphql.FingerprintQueryInput
// map fingerprints to their scene index
fpToScene := make(map[string][]int)
@@ -93,18 +93,27 @@ func (c Client) FindStashBoxScenesByFingerprints(sceneIDs []string) ([][]*models
}
if scene.Checksum.Valid {
fingerprints = append(fingerprints, scene.Checksum.String)
fingerprints = append(fingerprints, &graphql.FingerprintQueryInput{
Hash: scene.Checksum.String,
Algorithm: graphql.FingerprintAlgorithmMd5,
})
fpToScene[scene.Checksum.String] = append(fpToScene[scene.Checksum.String], index)
}
if scene.OSHash.Valid {
fingerprints = append(fingerprints, scene.OSHash.String)
fingerprints = append(fingerprints, &graphql.FingerprintQueryInput{
Hash: scene.OSHash.String,
Algorithm: graphql.FingerprintAlgorithmOshash,
})
fpToScene[scene.OSHash.String] = append(fpToScene[scene.OSHash.String], index)
}
if scene.Phash.Valid {
phashStr := utils.PhashToString(scene.Phash.Int64)
fingerprints = append(fingerprints, phashStr)
fingerprints = append(fingerprints, &graphql.FingerprintQueryInput{
Hash: phashStr,
Algorithm: graphql.FingerprintAlgorithmPhash,
})
fpToScene[phashStr] = append(fpToScene[phashStr], index)
}
}
@@ -147,7 +156,7 @@ func (c Client) FindStashBoxScenesByFingerprintsFlat(sceneIDs []string) ([]*mode
return nil, err
}
var fingerprints []string
var fingerprints []*graphql.FingerprintQueryInput
if err := c.txnManager.WithReadTxn(ctx, func(r models.ReaderRepository) error {
qb := r.Scene()
@@ -163,16 +172,24 @@ func (c Client) FindStashBoxScenesByFingerprintsFlat(sceneIDs []string) ([]*mode
}
if scene.Checksum.Valid {
fingerprints = append(fingerprints, scene.Checksum.String)
fingerprints = append(fingerprints, &graphql.FingerprintQueryInput{
Hash: scene.Checksum.String,
Algorithm: graphql.FingerprintAlgorithmMd5,
})
}
if scene.OSHash.Valid {
fingerprints = append(fingerprints, scene.OSHash.String)
fingerprints = append(fingerprints, &graphql.FingerprintQueryInput{
Hash: scene.OSHash.String,
Algorithm: graphql.FingerprintAlgorithmOshash,
})
}
if scene.Phash.Valid {
phashStr := utils.PhashToString(scene.Phash.Int64)
fingerprints = append(fingerprints, phashStr)
fingerprints = append(fingerprints, &graphql.FingerprintQueryInput{
Hash: utils.PhashToString(scene.Phash.Int64),
Algorithm: graphql.FingerprintAlgorithmPhash,
})
}
}
@@ -184,20 +201,20 @@ func (c Client) FindStashBoxScenesByFingerprintsFlat(sceneIDs []string) ([]*mode
return c.findStashBoxScenesByFingerprints(ctx, fingerprints)
}
func (c Client) findStashBoxScenesByFingerprints(ctx context.Context, fingerprints []string) ([]*models.ScrapedScene, error) {
func (c Client) findStashBoxScenesByFingerprints(ctx context.Context, fingerprints []*graphql.FingerprintQueryInput) ([]*models.ScrapedScene, error) {
var ret []*models.ScrapedScene
for i := 0; i < len(fingerprints); i += 100 {
end := i + 100
if end > len(fingerprints) {
end = len(fingerprints)
}
scenes, err := c.client.FindScenesByFingerprints(ctx, fingerprints[i:end])
scenes, err := c.client.FindScenesByFullFingerprints(ctx, fingerprints[i:end])
if err != nil {
return nil, err
}
sceneFragments := scenes.FindScenesByFingerprints
sceneFragments := scenes.FindScenesByFullFingerprints
for _, s := range sceneFragments {
ss, err := sceneFragmentToScrapedScene(ctx, c.getHTTPClient(), c.txnManager, s)