Add resolution extraction from JPEG screenshots in tester

This commit is contained in:
eduard256
2026-03-26 10:40:41 +00:00
parent 74b4b61198
commit 8dc8ba1096
3 changed files with 45 additions and 19 deletions
+21 -3
View File
@@ -70,7 +70,7 @@ func rtspHandler(rawURL string) (core.Producer, error) {
}
```
**Data flow**: URL -> GetHandler(url) -> handler(url) -> core.Producer -> GetMedias() -> codecs, latency -> getScreenshot() -> Result
**Data flow**: URL -> GetHandler(url) -> handler(url) -> core.Producer -> GetMedias() -> codecs, latency -> getScreenshot() -> jpegSize() -> Result (with width, height)
**Key**: The handler ONLY needs to return a `core.Producer`. Everything else (codecs extraction, screenshot capture, session management) is handled automatically by `worker.go`.
@@ -491,7 +491,7 @@ The tester uses:
2. `GetTrack()` + `Start()` -- to capture screenshot (keyframe)
3. `Stop()` -- to clean up
### How screenshot works (pkg/tester/worker.go)
### How screenshot and resolution work (pkg/tester/worker.go)
1. `getScreenshot(prod)` is called after successful Dial
2. Creates `magic.NewKeyframe()` consumer
@@ -501,8 +501,26 @@ The tester uses:
6. Waits for first keyframe via `cons.WriteTo()` with 10s timeout
7. If H264/H265 -- converts to JPEG via ffmpeg
8. If already JPEG -- uses as-is
9. `jpegSize(jpeg)` extracts width and height from JPEG SOF0/SOF2 marker
10. Resolution stored in `Result.Width` and `Result.Height`
This works automatically for ANY protocol that returns a valid `core.Producer`. You do NOT need to implement screenshot logic per protocol.
This works automatically for ANY protocol that returns a valid `core.Producer`. You do NOT need to implement screenshot or resolution logic per protocol.
### Result struct (pkg/tester/session.go)
```go
type Result struct {
Source string `json:"source"`
Screenshot string `json:"screenshot,omitempty"`
Codecs []string `json:"codecs,omitempty"`
Width int `json:"width,omitempty"` // from JPEG screenshot
Height int `json:"height,omitempty"` // from JPEG screenshot
LatencyMs int64 `json:"latency_ms,omitempty"`
Skipped bool `json:"skipped,omitempty"`
}
```
Resolution is extracted from the JPEG screenshot, not from SDP or protocol-specific data. This means width/height are only available when a screenshot was successfully captured. The frontend uses these values to classify streams as Main (HD) or Sub (SD).
### magic.NewKeyframe() (pkg/magic/keyframe.go)
+2 -2
View File
@@ -27,8 +27,8 @@ type Result struct {
Source string `json:"source"`
Screenshot string `json:"screenshot,omitempty"`
Codecs []string `json:"codecs,omitempty"`
Width uint16 `json:"width,omitempty"`
Height uint16 `json:"height,omitempty"`
Width int `json:"width,omitempty"`
Height int `json:"height,omitempty"`
LatencyMs int64 `json:"latency_ms,omitempty"`
Skipped bool `json:"skipped,omitempty"`
}
+22 -14
View File
@@ -7,7 +7,6 @@ import (
"time"
"github.com/AlexxIT/go2rtc/pkg/core"
"github.com/AlexxIT/go2rtc/pkg/h264"
"github.com/AlexxIT/go2rtc/pkg/magic"
)
@@ -67,31 +66,18 @@ func testURL(s *Session, rawURL string) {
latency := time.Since(start).Milliseconds()
var codecs []string
var width, height uint16
for _, media := range prod.GetMedias() {
if media.Direction != core.DirectionRecvonly {
continue
}
for _, codec := range media.Codecs {
codecs = append(codecs, codec.Name)
// extract resolution from first video codec SPS
if width == 0 && codec.Name == core.CodecH264 {
if spsBytes, _ := h264.GetParameterSet(codec.FmtpLine); spsBytes != nil {
if sps := h264.DecodeSPS(spsBytes); sps != nil {
width, height = sps.Width(), sps.Height()
}
}
}
}
}
r := &Result{
Source: rawURL,
Codecs: codecs,
Width: width,
Height: height,
LatencyMs: latency,
}
@@ -110,6 +96,7 @@ func testURL(s *Session, rawURL string) {
if jpeg != nil {
idx := s.AddScreenshot(jpeg)
r.Screenshot = fmt.Sprintf("/api/test/screenshot?id=%s&i=%d", s.ID, idx)
r.Width, r.Height = jpegSize(jpeg)
}
}
@@ -167,6 +154,27 @@ matched:
return once.Buffer(), cons.CodecName()
}
// jpegSize extracts width and height from JPEG SOF0/SOF2 marker
func jpegSize(data []byte) (int, int) {
for i := 2; i < len(data)-9; {
if data[i] != 0xFF {
return 0, 0
}
marker := data[i+1]
size := int(data[i+2])<<8 | int(data[i+3])
// SOF0 (0xC0) or SOF2 (0xC2) -- baseline or progressive
if marker == 0xC0 || marker == 0xC2 {
h := int(data[i+5])<<8 | int(data[i+6])
w := int(data[i+7])<<8 | int(data[i+8])
return w, h
}
i += 2 + size
}
return 0, 0
}
func toJPEG(raw []byte) []byte {
cmd := exec.Command("ffmpeg",
"-hide_banner", "-loglevel", "error",