vphash/vphash.go

package vphash

import (
	"image"
	"time"

	"github.com/corona10/goimagehash"
)

// Decoder is the minimal interface VPHash needs.
// Implement this for any video source: MPEG, AV1, H.264, live stream.
type Decoder interface {
	Next() (img image.Image, t time.Duration, ok bool)
}

type Entry struct {
	Start time.Duration
	End   time.Duration
	Hash  *goimagehash.ImageHash
}

// Scan reads frames and emits scene entries as they're detected.
// The channel closes when the video ends.
func Scan(dec Decoder, maxDistance, minFrames int) <-chan Entry {
	ch := make(chan Entry)

	if maxDistance == 0 {
		maxDistance = 10
	}

	if minFrames == 0 {
		minFrames = 24
	}

	go func() {
		defer close(ch)

		var (
			hashes     []*goimagehash.ImageHash
			times      []time.Duration
			sceneStart int
		)

		for {
			img, t, ok := dec.Next()
			if !ok {
				if len(hashes)-sceneStart >= minFrames {
					mid := (sceneStart + len(hashes) - 1) / 2
					ch <- Entry{Start: times[sceneStart], End: t, Hash: hashes[mid]}
				}
				return
			}

			h, _ := goimagehash.PerceptionHash(img)
			hashes = append(hashes, h)
			times = append(times, t)

			if len(hashes) == 1 {
				continue
			}

			dist, _ := hashes[len(hashes)-2].Distance(h)
			if dist > maxDistance {
				if len(hashes)-1-sceneStart >= minFrames {
					mid := (sceneStart + len(hashes) - 2) / 2
					ch <- Entry{Start: times[sceneStart], End: t, Hash: hashes[mid]}
				}
				sceneStart = len(hashes) - 1
			}
		}
	}()

	return ch
}