Merge pull request #3 from acolombier/feat/add-cover-support

feat: add support for MP4 and ID3v2 cover

.github/ISSUE_TEMPLATE/bug.yaml

@@ -19,7 +19,7 @@ body:
     attributes:
       label: Version
       description: What version of stemgen are you running?
-      placeholder: 0.1.0, main, ...
+      placeholder: 0.2.0, main, ...
   - type: input
     id: os
     attributes:

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.10 as build
+FROM python:3.11 as build
 WORKDIR /build
 RUN wget -O taglib.tar.gz https://github.com/taglib/taglib/releases/download/v2.0.1/taglib-2.0.1.tar.gz && \
     tar xf taglib.tar.gz  && \
@@ -12,13 +12,16 @@ RUN python3 -m build --wheel && \
     find /usr/lib -name libtag.so.2.0.1 -exec cp '{}' /build/libtag.so.2.0.1 \;
 
 
-FROM python:3.10
+FROM python:3.11
 COPY --from=build /build/libtag.so.2.0.1 /usr/local/lib/libtag.so.2.0.1
 COPY --from=build /build/dist/stemgen-*.whl /tmp/
 COPY --from=build /usr/include/taglib/ /usr/include/taglib/
 RUN ln -s /usr/local/lib/libtag.so.2.0.1 /usr/local/lib/libtag.so && \
+    apt update && apt install -y ffmpeg libboost-python1.74-dev libboost-python1.74.0 && \
     python -m pip install /tmp/stemgen-*.whl && \
-    apt update && apt install -y ffmpeg && \
     pip install --upgrade --force torchaudio && \
-    rm -rf /root/.cache /tmp/* && rm -rf /usr/include/taglib/
+    apt-get purge -y libboost-python1.74-dev && \
+    apt-get clean autoclean && \
+    apt-get autoremove --yes && \
+    rm -rf /root/.cache /tmp/* && rm -rf /usr/include/taglib/ && rm -rf /var/lib/{apt,dpkg,cache,log}/
 CMD ["/usr/local/bin/stemgen"]

README.md

@@ -36,7 +36,7 @@ Under the hood, it uses:
 > isn't supported
 
 ```sh
-pip install -e "git+https://github.com/acolombier/stemgen.git@0.1.0#egg=stemgen"
+pip install -e "git+https://github.com/acolombier/stemgen.git@0.2.0#egg=stemgen"
 ```
 
 ### Ubuntu 22.04 / Debian Bookworm / PopOS 22.04
@@ -191,7 +191,7 @@ container. Here the simple way to use it:
 docker run \
     -v /path/to/folder:/path/to/folder \
     -it --rm \
-    aclmb/stemgen:0.1.0 \
+    aclmb/stemgen:0.2.0 \
         /path/to/folder/Artist\ -\ Title.mp3 \
         /path/to/folder
 ```
@@ -204,7 +204,7 @@ docker run \
     -v /path/to/folder:/path/to/folder \
     -v stemgen_torch_cache:/root/.cache/torch/hub/ \
     -it --gpus --rm \
-    aclmb/stemgen:0.1.0 \
+    aclmb/stemgen:0.2.0 \
         /path/to/folder/Artist\ -\ Title.mp3 \
         /path/to/folder
 ```

pyproject.toml

@@ -1,3 +1,7 @@
 [build-system]
 requires = ["setuptools>=61.0", "Cython>=3.0.10"]
 build-backend = "setuptools.build_meta"
+
+
+[tool.licensecheck]
+using = "requirements:requirements.txt"

requirements.txt

@@ -1,6 +1,6 @@
 demucs @ git+https://github.com/facebookresearch/demucs.git@583db9df0213ba5f5b3491eca5c993e7629f1949#egg=demucs
+tagpy @ git+https://github.com/acolombier/tagpy.git@c5de51fe9636b312bfe95dee8b051ab640976d43#egg=tagpy
 setuptools>=61.0
-pytaglib==3.0.0
 ffmpeg-python==0.2.0
 torch>=2.1.2
 torchaudio>=2.1.2

setup.py

@@ -13,7 +13,7 @@
 URL = "https://github.com/acolombier/stemgen"
 EMAIL = "[email protected]"
 AUTHOR = "Antoine Colombier"
-REQUIRES_PYTHON = ">=3.10.0"
+REQUIRES_PYTHON = ">=3.11.0"
 
 # Get version without explicitly loading the module.
 for line in open("stemgen/__init__.py"):

stemgen/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.0"
+__version__ = "0.2.0"

stemgen/nistemfile.py

@@ -1,12 +1,56 @@
 import click
 
-import taglib
+import tagpy
+import tagpy.id3v2
 from torchaudio.io import StreamWriter, CodecConfig
 import stembox
 import torch
 
 from .constant import SAMPLE_RATE, CHUNK_SIZE
 
+SUPPORTED_TAGS = [
+    "title",
+    "artist",
+    "album",
+    "comment",
+    "genre",
+    "year",
+    "track",
+]
+
+
+def _extract_cover(f):
+    tag = None
+    if isinstance(f, tagpy.FileRef):
+        tag = f.tag()
+        f = f.file()
+    covers = []
+    if hasattr(tag, "covers"):
+        covers = tag.covers
+    elif hasattr(f, "ID3v2Tag"):
+        covers = [
+            a
+            for a in f.ID3v2Tag().frameList()
+            if isinstance(a, tagpy.id3v2.AttachedPictureFrame)
+        ]
+
+    if covers:
+        cover = covers[0]
+        fmt = tagpy.mp4.CoverArtFormats.Unknown
+        if isinstance(cover, tagpy.mp4.CoverArt):
+            return cover
+        else:
+            mime = cover.mimeType().lower().strip()
+            if "image/jpeg":
+                fmt = tagpy.mp4.CoverArtFormats.JPEG
+            elif "image/png":
+                fmt = tagpy.mp4.CoverArtFormats.PNG
+            elif "image/bmp":
+                fmt = tagpy.mp4.CoverArtFormats.BMP
+            elif "image/gif":
+                fmt = tagpy.mp4.CoverArtFormats.GIF
+            return tagpy.mp4.CoverArt(fmt, cover.picture())
+
 
 class NIStemFile:
     STEM_DEFAULT_LABEL = [
@@ -74,10 +118,20 @@ def write(self, original, stems):
                 progress.finish()
 
     def update_metadata(self, src, **stem_metadata):
-        with taglib.File(src) as src, taglib.File(
-            self.__path, save_on_exit=True
-        ) as dst:
-            dst.tags = src.tags
+        src = tagpy.FileRef(src)
+        dst = tagpy.FileRef(self.__path)
+
+        src_tag = src.tag()
+        dst_tag = dst.tag()
+        for tag in SUPPORTED_TAGS:
+            setattr(dst_tag, tag, getattr(src_tag, tag))
+
+        cover = _extract_cover(src)
+        if cover:
+            c = tagpy.mp4.CoverArtList()
+            c.append(cover)
+            dst_tag.covers = c
+        dst.save()
 
         with stembox.Stem(self.__path) as f:
             f.stems = [