Add ranges to funscript heatmaps (#3373)

* Add range to funscript heatmaps
* Add config to draw heatmap range
* Add setting to UI
* Apply draw range setting

Includes some refactoring

---------

Co-authored-by: kermieisinthehouse <kermie@isinthe.house>
Co-authored-by: WithoutPants <53250216+WithoutPants@users.noreply.github.com>

graphql/documents/data/config.graphql

@@ -51,6 +51,7 @@ fragment ConfigGeneralData on ConfigGeneralResult {
   transcodeOutputArgs
   liveTranscodeInputArgs
   liveTranscodeOutputArgs
+  drawFunscriptHeatmapRange
 }
 
 fragment ConfigInterfaceData on ConfigInterfaceResult {

graphql/schema/types/config.graphql

@@ -84,6 +84,9 @@ input ConfigGeneralInput {
   These are applied when live transcoding"""
   liveTranscodeOutputArgs: [String!]
 
+  """whether to include range in generated funscript heatmaps"""
+  drawFunscriptHeatmapRange: Boolean
+
   """Write image thumbnails to disk when generating on the fly"""
   writeImageThumbnails: Boolean
   """Username"""
@@ -184,6 +187,9 @@ type ConfigGeneralResult {
   These are applied when live transcoding"""
   liveTranscodeOutputArgs: [String!]!
 
+  """whether to include range in generated funscript heatmaps"""
+  drawFunscriptHeatmapRange: Boolean!
+
   """Write image thumbnails to disk when generating on the fly"""
   writeImageThumbnails: Boolean!
   """API Key"""

internal/api/resolver_mutation_configure.go

@@ -316,6 +316,10 @@ func (r *mutationResolver) ConfigureGeneral(ctx context.Context, input ConfigGen
 		c.Set(config.LiveTranscodeOutputArgs, input.LiveTranscodeOutputArgs)
 	}
 
+	if input.DrawFunscriptHeatmapRange != nil {
+		c.Set(config.DrawFunscriptHeatmapRange, input.DrawFunscriptHeatmapRange)
+	}
+
 	if err := c.Write(); err != nil {
 		return makeConfigGeneralResult(), err
 	}

internal/api/resolver_query_configuration.go

@@ -128,6 +128,7 @@ func makeConfigGeneralResult() *ConfigGeneralResult {
 		TranscodeOutputArgs:          config.GetTranscodeOutputArgs(),
 		LiveTranscodeInputArgs:       config.GetLiveTranscodeInputArgs(),
 		LiveTranscodeOutputArgs:      config.GetLiveTranscodeOutputArgs(),
+		DrawFunscriptHeatmapRange:    config.GetDrawFunscriptHeatmapRange(),
 	}
 }
 

internal/manager/config/config.go

@@ -185,6 +185,9 @@ const (
 	HandyKey        = "handy_key"
 	FunscriptOffset = "funscript_offset"
 
+	DrawFunscriptHeatmapRange        = "draw_funscript_heatmap_range"
+	drawFunscriptHeatmapRangeDefault = true
+
 	ThemeColor        = "theme_color"
 	DefaultThemeColor = "#202b33"
 
@@ -831,6 +834,10 @@ func (i *Instance) GetLiveTranscodeOutputArgs() []string {
 	return i.getStringSlice(LiveTranscodeOutputArgs)
 }
 
+func (i *Instance) GetDrawFunscriptHeatmapRange() bool {
+	return i.getBoolDefault(DrawFunscriptHeatmapRange, drawFunscriptHeatmapRangeDefault)
+}
+
 // IsWriteImageThumbnails returns true if image thumbnails should be written
 // to disk after generating on the fly.
 func (i *Instance) IsWriteImageThumbnails() bool {

internal/manager/generator_interactive_heatmap_speed.go

@@ -15,14 +15,13 @@ import (
 )
 
 type InteractiveHeatmapSpeedGenerator struct {
-	sceneDurationMilli int64
-	InteractiveSpeed   int
-	Funscript          Script
-	FunscriptPath      string
-	HeatmapPath        string
-	Width              int
-	Height             int
-	NumSegments        int
+	InteractiveSpeed int
+	Funscript        Script
+	Width            int
+	Height           int
+	NumSegments      int
+
+	DrawRange bool
 }
 
 type Script struct {
@@ -33,8 +32,7 @@ type Script struct {
 	// Range is the percentage of a full stroke to use.
 	Range int `json:"range,omitempty"`
 	// Actions are the timed moves.
-	Actions      []Action `json:"actions"`
-	AvarageSpeed int64
+	Actions []Action `json:"actions"`
 }
 
 // Action is a move at a specific time.
@@ -50,23 +48,22 @@ type Action struct {
 }
 
 type GradientTable []struct {
-	Col colorful.Color
-	Pos float64
+	Col    colorful.Color
+	Pos    float64
+	YRange [2]float64
 }
 
-func NewInteractiveHeatmapSpeedGenerator(funscriptPath string, heatmapPath string, sceneDuration float64) *InteractiveHeatmapSpeedGenerator {
+func NewInteractiveHeatmapSpeedGenerator(drawRange bool) *InteractiveHeatmapSpeedGenerator {
 	return &InteractiveHeatmapSpeedGenerator{
-		sceneDurationMilli: int64(sceneDuration * 1000),
-		FunscriptPath:      funscriptPath,
-		HeatmapPath:        heatmapPath,
-		Width:              320,
-		Height:             15,
-		NumSegments:        150,
+		Width:       1280,
+		Height:      60,
+		NumSegments: 600,
+		DrawRange:   drawRange,
 	}
 }
 
-func (g *InteractiveHeatmapSpeedGenerator) Generate() error {
-	funscript, err := g.LoadFunscriptData(g.FunscriptPath)
+func (g *InteractiveHeatmapSpeedGenerator) Generate(funscriptPath string, heatmapPath string, sceneDuration float64) error {
+	funscript, err := g.LoadFunscriptData(funscriptPath, sceneDuration)
 
 	if err != nil {
 		return err
@@ -79,7 +76,7 @@ func (g *InteractiveHeatmapSpeedGenerator) Generate() error {
 	g.Funscript = funscript
 	g.Funscript.UpdateIntensityAndSpeed()
 
-	err = g.RenderHeatmap()
+	err = g.RenderHeatmap(heatmapPath)
 
 	if err != nil {
 		return err
@@ -90,7 +87,7 @@ func (g *InteractiveHeatmapSpeedGenerator) Generate() error {
 	return nil
 }
 
-func (g *InteractiveHeatmapSpeedGenerator) LoadFunscriptData(path string) (Script, error) {
+func (g *InteractiveHeatmapSpeedGenerator) LoadFunscriptData(path string, sceneDuration float64) (Script, error) {
 	data, err := os.ReadFile(path)
 	if err != nil {
 		return Script{}, err
@@ -111,8 +108,9 @@ func (g *InteractiveHeatmapSpeedGenerator) LoadFunscriptData(path string) (Scrip
 	// trim actions with negative timestamps to avoid index range errors when generating heatmap
 	// #3181 - also trim actions that occur after the scene duration
 	loggedBadTimestamp := false
+	sceneDurationMilli := int64(sceneDuration * 1000)
 	isValid := func(x int64) bool {
-		return x >= 0 && x < g.sceneDurationMilli
+		return x >= 0 && x < sceneDurationMilli
 	}
 
 	i := 0
@@ -157,14 +155,27 @@ func (funscript *Script) UpdateIntensityAndSpeed() {
 }
 
 // funscript needs to have intensity updated first
-func (g *InteractiveHeatmapSpeedGenerator) RenderHeatmap() error {
-
+func (g *InteractiveHeatmapSpeedGenerator) RenderHeatmap(heatmapPath string) error {
 	gradient := g.Funscript.getGradientTable(g.NumSegments)
 
 	img := image.NewRGBA(image.Rect(0, 0, g.Width, g.Height))
 	for x := 0; x < g.Width; x++ {
-		c := gradient.GetInterpolatedColorFor(float64(x) / float64(g.Width))
-		draw.Draw(img, image.Rect(x, 0, x+1, g.Height), &image.Uniform{c}, image.Point{}, draw.Src)
+		xPos := float64(x) / float64(g.Width)
+		c := gradient.GetInterpolatedColorFor(xPos)
+
+		y0 := 0
+		y1 := g.Height
+
+		if g.DrawRange {
+			yRange := gradient.GetYRange(xPos)
+			top := int(yRange[0] / 100.0 * float64(g.Height))
+			bottom := int(yRange[1] / 100.0 * float64(g.Height))
+
+			y0 = g.Height - top
+			y1 = g.Height - bottom
+		}
+
+		draw.Draw(img, image.Rect(x, y0, x+1, y1), &image.Uniform{c}, image.Point{}, draw.Src)
 	}
 
 	// add 10 minute marks
@@ -178,7 +189,7 @@ func (g *InteractiveHeatmapSpeedGenerator) RenderHeatmap() error {
 		ts += tick
 	}
 
-	outpng, err := os.Create(g.HeatmapPath)
+	outpng, err := os.Create(heatmapPath)
 	if err != nil {
 		return err
 	}
@@ -217,27 +228,96 @@ func (gt GradientTable) GetInterpolatedColorFor(t float64) colorful.Color {
 	return gt[len(gt)-1].Col
 }
 
+func (gt GradientTable) GetYRange(t float64) [2]float64 {
+	for i := 0; i < len(gt)-1; i++ {
+		c1 := gt[i]
+		c2 := gt[i+1]
+		if c1.Pos <= t && t <= c2.Pos {
+			// TODO: We are in between c1 and c2. Go blend them!
+			return c1.YRange
+		}
+	}
+
+	// Nothing found? Means we're at (or past) the last gradient keypoint.
+	return gt[len(gt)-1].YRange
+}
+
 func (funscript Script) getGradientTable(numSegments int) GradientTable {
+	const windowSize = 15
+	const backfillThreshold = 500
+
 	segments := make([]struct {
 		count     int
 		intensity int
+		yRange    [2]float64
+		at        int64
 	}, numSegments)
 	gradient := make(GradientTable, numSegments)
+	posList := []int{}
 
 	maxts := funscript.Actions[len(funscript.Actions)-1].At
 
 	for _, a := range funscript.Actions {
+		posList = append(posList, a.Pos)
+
+		if len(posList) > windowSize {
+			posList = posList[1:]
+		}
+
+		sortedPos := make([]int, len(posList))
+		copy(sortedPos, posList)
+		sort.Ints(sortedPos)
+
+		topHalf := sortedPos[len(sortedPos)/2:]
+		bottomHalf := sortedPos[0 : len(sortedPos)/2]
+
+		var totalBottom int
+		var totalTop int
+
+		for _, value := range bottomHalf {
+			totalBottom += value
+		}
+		for _, value := range topHalf {
+			totalTop += value
+		}
+
+		averageBottom := float64(totalBottom) / float64(len(bottomHalf))
+		averageTop := float64(totalTop) / float64(len(topHalf))
+
 		segment := int(float64(a.At) / float64(maxts+1) * float64(numSegments))
 		// #3181 - sanity check. Clamp segment to numSegments-1
 		if segment >= numSegments {
 			segment = numSegments - 1
 		}
+		segments[segment].at = a.At
 		segments[segment].count++
 		segments[segment].intensity += int(a.Intensity)
+		segments[segment].yRange[0] = averageTop
+		segments[segment].yRange[1] = averageBottom
+	}
+
+	lastSegment := segments[0]
+
+	// Fill in gaps in segments
+	for i := 0; i < numSegments; i++ {
+		segmentTS := int64(float64(i) / float64(numSegments))
+
+		// Empty segment - fill it with the previous up to backfillThreshold ms
+		if segments[i].count == 0 {
+			if segmentTS-lastSegment.at < backfillThreshold {
+				segments[i].count = lastSegment.count
+				segments[i].intensity = lastSegment.intensity
+				segments[i].yRange[0] = lastSegment.yRange[0]
+				segments[i].yRange[1] = lastSegment.yRange[1]
+			}
+		} else {
+			lastSegment = segments[i]
+		}
 	}
 
 	for i := 0; i < numSegments; i++ {
 		gradient[i].Pos = float64(i) / float64(numSegments-1)
+		gradient[i].YRange = segments[i].yRange
 		if segments[i].count > 0 {
 			gradient[i].Col = getSegmentColor(float64(segments[i].intensity) / float64(segments[i].count))
 		} else {

internal/manager/task_generate_interactive_heatmap_speed.go

@@ -29,10 +29,11 @@ func (t *GenerateInteractiveHeatmapSpeedTask) Start(ctx context.Context) {
 	videoChecksum := t.Scene.GetHash(t.fileNamingAlgorithm)
 	funscriptPath := video.GetFunscriptPath(t.Scene.Path)
 	heatmapPath := instance.Paths.Scene.GetInteractiveHeatmapPath(videoChecksum)
+	drawRange := instance.Config.GetDrawFunscriptHeatmapRange()
 
-	generator := NewInteractiveHeatmapSpeedGenerator(funscriptPath, heatmapPath, t.Scene.Files.Primary().Duration)
+	generator := NewInteractiveHeatmapSpeedGenerator(drawRange)
 
-	err := generator.Generate()
+	err := generator.Generate(funscriptPath, heatmapPath, t.Scene.Files.Primary().Duration)
 
 	if err != nil {
 		logger.Errorf("error generating heatmap: %s", err.Error())

ui/v2.5/src/components/Settings/SettingsSystemPanel.tsx

@@ -315,6 +315,16 @@ export const SettingsConfigurationPanel: React.FC = () => {
         />
       </SettingSection>
 
+      <SettingSection headingID="config.general.heatmap_generation">
+        <BooleanSetting
+          id="heatmap-draw-range"
+          headingID="config.general.funscript_heatmap_draw_range"
+          subHeadingID="config.general.funscript_heatmap_draw_range_desc"
+          checked={general.drawFunscriptHeatmapRange ?? true}
+          onChange={(v) => saveGeneral({ drawFunscriptHeatmapRange: v })}
+        />
+      </SettingSection>
+
       <SettingSection headingID="config.general.logging">
         <StringSetting
           headingID="config.general.auth.log_file"

ui/v2.5/src/docs/en/Changelog/v0200.md

@@ -1,4 +1,5 @@
 ### ✨ New Features
+* Optionally show range in generated funscript heatmaps. ([#3373](https://github.com/stashapp/stash/pull/3373))
 * Show funscript heatmaps in scene player scrubber. ([#3374](https://github.com/stashapp/stash/pull/3374))
 * Support customising the filename regex used for determining the gallery cover image. ([#3391](https://github.com/stashapp/stash/pull/3391))
 * Added tenth-place rating precision option. ([#3343](https://github.com/stashapp/stash/pull/3343))

ui/v2.5/src/locales/en-GB.json

@@ -296,6 +296,8 @@
           }
         }
       },
+      "funscript_heatmap_draw_range": "Include range in generated heatmaps",
+      "funscript_heatmap_draw_range_desc": "Draw range of motion on the y-axis of the generated heatmap. Existing heatmaps will need to be regenerated after changing.",
       "gallery_ext_desc": "Comma-delimited list of file extensions that will be identified as gallery zip files.",
       "gallery_ext_head": "Gallery zip Extensions",
       "generated_file_naming_hash_desc": "Use MD5 or oshash for generated file naming. Changing this requires that all scenes have the applicable MD5/oshash value populated. After changing this value, existing generated files will need to be migrated or regenerated. See Tasks page for migration.",
@@ -303,6 +305,7 @@
       "generated_files_location": "Directory location for the generated files (scene markers, scene previews, sprites, etc)",
       "generated_path_head": "Generated Path",
       "hashing": "Hashing",
+      "heatmap_generation": "Funscript Heatmap Generation",
       "image_ext_desc": "Comma-delimited list of file extensions that will be identified as images.",
       "image_ext_head": "Image Extensions",
       "include_audio_desc": "Includes audio stream when generating previews.",