Convert to three-channel if needed in spectrogram_from_image

Closes: #51 Topic: convert_three_channel
2022-12-29 11:27:42 -08:00 · 2022-12-29 11:27:42 -08:00 · ce8744897b
parent 89f3946155
commit ce8744897b
1 changed files with 4 additions and 0 deletions
--- a/riffusion/util/image_util.py
+++ b/riffusion/util/image_util.py
@ -77,6 +77,10 @@ def spectrogram_from_image(
    Returns:
        spectrogram: (channels, frequency, time)
    """
+    # Convert to RGB if single channel
+    if image.mode in ("P", "L"):
+        image = image.convert("RGB")
+
    # Flip Y
    image = image.transpose(Image.Transpose.FLIP_TOP_BOTTOM)