Process images when they're of YCbCr type

viamrobotics · Oct 5, 2022 · 29c1e15 · 29c1e15
1 parent 52cee9f
commit 29c1e15
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/services/slam/builtin/builtin.go b/services/slam/builtin/builtin.go
@@ -736,13 +736,22 @@ func (slamSvc *builtIn) getLazyPNGImage(ctx context.Context, cam camera.Camera)
 	}
 
 	lazyImg, ok := img.(*rimage.LazyEncodedImage)
-	if !ok {
-		return nil, errors.Errorf("expected lazily encoded image, got %T", lazyImg)
+	if ok {
+		if lazyImg.MIMEType() != utils.MimeTypePNG {
+			return nil, errors.Errorf("expected mime type %v, got %T", utils.MimeTypePNG, img)
+		}
+		return lazyImg.RawData(), nil
 	}
-	if lazyImg.MIMEType() != utils.MimeTypePNG {
-		return nil, errors.Errorf("expected mime type %v, got %v", utils.MimeTypePNG, lazyImg.MIMEType())
+	ycbcrImg, ok := img.(*image.YCbCr)
+	if ok {
+		pngImage, err := rimage.EncodeImage(ctx, ycbcrImg, utils.MimeTypePNG)
+		if err != nil {
+			return nil, err
+		}
+		return pngImage, nil
 	}
-	return lazyImg.RawData(), nil
+
+	return nil, errors.Errorf("expected lazily encoded image or ycbcrImg, got %T", img)
 }
 
 // getAndSaveDataSparse implements the data extraction for sparse algos and saving to the directory path (data subfolder) specified in