diff --git a/vision/objectdetection/detection_source.go b/vision/objectdetection/detection_source.go index 35b67534574..8d4f0b8e733 100644 --- a/vision/objectdetection/detection_source.go +++ b/vision/objectdetection/detection_source.go @@ -32,6 +32,9 @@ type Source struct { cancelFunc func() } +// WaitForMs tells us the number of milliseconds to wait on the channel for. +const WaitForMs = 2000 + // NewSource builds the pipeline from an input VideoSource and Detector. func NewSource(src gostream.VideoSource, det Detector) (*Source, error) { // fill optional functions with identity operators @@ -75,6 +78,7 @@ func (s *Source) backgroundWorker(stream gostream.VideoStream, det Detector) { Release: release, Err: err, } + select { case <-s.cancelCtx.Done(): return @@ -96,10 +100,12 @@ func (s *Source) Read(ctx context.Context) (image.Image, func(), error) { ctx, span := trace.StartSpan(ctx, "vision::objectdetection::Source::Read") defer span.End() start := time.Now() + res, err := s.NextResult(ctx) if err != nil { return nil, nil, err } + duration := time.Since(start) fps := 1. / duration.Seconds() ovImg, err := Overlay(res.OriginalImage, res.Detections) @@ -114,6 +120,7 @@ func (s *Source) Read(ctx context.Context) (image.Image, func(), error) { func (s *Source) NextResult(ctx context.Context) (*Result, error) { ctx, span := trace.StartSpan(ctx, "vision::objectdetection::Source::NextResult") defer span.End() + select { case <-ctx.Done(): return nil, ctx.Err() @@ -121,5 +128,7 @@ func (s *Source) NextResult(ctx context.Context) (*Result, error) { return nil, s.cancelCtx.Err() case result := <-s.pipelineOutput: return result, result.Err + case <-time.After(WaitForMs * time.Millisecond): + return nil, errors.Errorf("nothing on channel after %v ms", WaitForMs) } } diff --git a/vision/objectdetection/detection_source_test.go b/vision/objectdetection/detection_source_test.go index d8972c11b5d..7a8d18ec013 100644 --- a/vision/objectdetection/detection_source_test.go +++ b/vision/objectdetection/detection_source_test.go @@ -4,6 +4,7 @@ import ( "context" "image" "testing" + "time" "go.viam.com/test" "go.viam.com/utils/artifact" @@ -20,6 +21,7 @@ func TestDetectionSource(t *testing.T) { test.That(t, err, test.ShouldBeNil) src, err := camera.NewFromReader(context.Background(), &videosource.StaticSource{ColorImg: sourceImg}, nil, camera.ColorStream) test.That(t, err, test.ShouldBeNil) + // make the preprocessing function p, err := objectdetection.RemoveColorChannel("b") test.That(t, err, test.ShouldBeNil) @@ -39,21 +41,37 @@ func TestDetectionSource(t *testing.T) { defer pipeline.Close() // compare with expected bounding boxes + start := time.Now() res, err := pipeline.NextResult(context.Background()) - test.That(t, err, test.ShouldBeNil) - bbs := res.Detections - test.That(t, bbs, test.ShouldHaveLength, 1) - test.That(t, bbs[0].Score(), test.ShouldEqual, 1.0) - test.That(t, bbs[0].Label(), test.ShouldEqual, "orange") - test.That(t, bbs[0].BoundingBox(), test.ShouldResemble, &image.Rectangle{image.Point{848, 424}, image.Point{999, 565}}) + tt := time.Now() + elapsed := tt.Sub(start) + if elapsed > objectdetection.WaitForMs*time.Millisecond { + test.That(t, err, test.ShouldNotBeNil) + test.That(t, res, test.ShouldBeNil) + } else { + test.That(t, err, test.ShouldBeNil) + bbs := res.Detections + test.That(t, bbs, test.ShouldHaveLength, 1) + test.That(t, bbs[0].Score(), test.ShouldEqual, 1.0) + test.That(t, bbs[0].Label(), test.ShouldEqual, "orange") + test.That(t, bbs[0].BoundingBox(), test.ShouldResemble, &image.Rectangle{image.Point{848, 424}, image.Point{999, 565}}) + } // overlay the image and see if it is red where you expect + start = time.Now() img, _, err := pipeline.Read(context.Background()) - test.That(t, err, test.ShouldBeNil) - ovImg := rimage.ConvertImage(img) - test.That(t, ovImg.GetXY(848, 424), test.ShouldResemble, rimage.Red) - test.That(t, ovImg.GetXY(998, 564), test.ShouldResemble, rimage.Red) - test.That(t, src.Close(context.Background()), test.ShouldBeNil) + tt = time.Now() + elapsed = tt.Sub(start) + if elapsed > objectdetection.WaitForMs*time.Millisecond { + test.That(t, err, test.ShouldNotBeNil) + test.That(t, img, test.ShouldBeNil) + } else { + test.That(t, err, test.ShouldBeNil) + ovImg := rimage.ConvertImage(img) + test.That(t, ovImg.GetXY(848, 424), test.ShouldResemble, rimage.Red) + test.That(t, ovImg.GetXY(998, 564), test.ShouldResemble, rimage.Red) + test.That(t, src.Close(context.Background()), test.ShouldBeNil) + } } func TestEmptyDetection(t *testing.T) {