media-toolkit 0.0.0__py3-none-any.whl

media_toolkit/__init__.py

@@ -0,0 +1 @@
+from media_toolkit.core import MediaFile, ImageFile, VideoFile, AudioFile

media_toolkit/core/__init__.py

@@ -0,0 +1,4 @@
+from .media_file import MediaFile
+from .image_file import ImageFile
+from .audio_file import AudioFile
+from media_toolkit.core.video.video_file import VideoFile

media_toolkit/core/audio_file.py

@@ -0,0 +1,38 @@
+from media_toolkit.dependency_requirements import requires
+from media_toolkit.core.media_file import MediaFile
+import io
+
+try:
+    import soundfile
+    import numpy as np
+except ImportError:
+    pass
+
+
+class AudioFile(MediaFile):
+    """
+    Has file conversions that make it easy to work with image files across the web.
+    Internally it uses numpy and librosa.
+    """
+    @requires('soundfile')
+    def to_soundfile(self):
+        return soundfile.read(self._content_buffer)
+
+    @requires('soundfile')
+    def to_np_array(self, sr: int = None, return_sample_rate: bool = False):
+        self._content_buffer.seek(0)
+        audio, sr = soundfile.read(self._content_buffer, samplerate=sr)  # sr=None returns the native sample rate
+        if return_sample_rate:
+            return audio, sr
+        return audio
+
+    @requires('soundfile')
+    def from_np_array(self, np_array, sr: int = None, file_type: str = "wav"):
+        sr = 22050 if sr is None else sr
+        # write to virtual file with librosa
+        virtual_file = io.BytesIO()
+        virtual_file.name = f"audio_file.{file_type}"
+        soundfile.write(virtual_file, np_array, samplerate=sr, format=file_type)
+
+        super().from_file(virtual_file)
+

media_toolkit/core/image_file.py

@@ -0,0 +1,89 @@
+from media_toolkit.dependency_requirements import requires_numpy, requires_cv2, requires
+from media_toolkit.core.media_file import MediaFile
+
+try:
+    import cv2
+    import numpy as np
+except ImportError:
+    pass
+
+
+class ImageFile(MediaFile):
+    """
+    Has file conversions that make it easy to work with image files across the web.
+    Internally it uses cv2 file format.
+    """
+    @requires('cv2', 'numpy')
+    def from_np_array(self, np_array, img_type: str = None):
+        if isinstance(np_array, list):
+            np_array = np.array(np_array)
+
+        if img_type is None:
+            if "image/" not in self.content_type:
+                img_type, self._channels = self.detect_image_type_and_channels(np_array)
+            else:
+                img_type = self.content_type.split("/")[1]
+            self.content_type = f"image/{img_type}"
+
+        is_success, buffer = cv2.imencode(f".{img_type}", np_array)
+        if is_success:
+            # avoid to check again for image type by calling super().from_bytes instead of self.from_bytes
+            return super().from_bytes(buffer)
+        else:
+            raise ValueError(f"Could not convert np_array to {img_type} image")
+
+
+    @requires('numpy', 'cv2')
+    def to_np_array(self):
+        bytes = self.to_bytes()
+        return cv2.imdecode(np.frombuffer(bytes, np.uint8), -1)
+
+    @requires_numpy()
+    def to_cv2_img(self):
+        return self.to_np_array()
+
+    @requires_cv2()
+    def save(self, path: str):
+        cv2.imwrite(path, self.to_np_array())
+
+    def _file_info(self):
+        super()._file_info()
+        np_array = self.to_np_array()
+        img_type, self._channels = self.detect_image_type_and_channels(np_array)
+        if img_type is not None:
+            self.content_type = f"image/{img_type}"
+        
+
+    @staticmethod
+    def detect_image_type_and_channels(image) -> (str, int):
+        """Detect the image type and number of _channels from a numpy array."""
+        if isinstance(image, list):
+            image = np.array(image)
+
+        # Check the number of _channels
+        if len(image.shape) == 2:
+            channels = 1  # Grayscale
+        elif len(image.shape) == 3:
+            channels = image.shape[2]
+        else:
+            #raise ValueError("Unsupported image shape: {}".format(image.shape))
+            return None, None
+
+        # Detect image type by checking for specific markers
+        image_type = None
+
+        # Convert to bytes and inspect file signatures for format detection
+        success, encoded_image = cv2.imencode('.png', image)
+        if success:
+            encoded_bytes = encoded_image.tobytes()
+            if encoded_bytes.startswith(b'\x89PNG\r\n\x1a\n'):
+                image_type = 'png'
+            elif encoded_bytes[0:2] == b'\xff\xd8':
+                image_type = 'jpeg'
+            elif encoded_bytes.startswith(b'BM'):
+                image_type = 'bmp'
+            elif encoded_bytes.startswith(b'GIF'):
+                image_type = 'gif'
+
+        return image_type, channels
+

media_toolkit/core/media_file.py

@@ -0,0 +1,259 @@
+import base64
+import io
+import mimetypes
+from typing import Union, BinaryIO
+import os
+
+from media_toolkit.dependency_requirements import requires_numpy
+
+
+try:
+    import numpy as np
+except ImportError:
+    pass
+
+
+class MediaFile:
+    """
+    Has file conversions that make it easy to work standardized with files across the web and in the sdk.
+    Works natively with bytesio, base64 and binary data.
+    """
+    def __init__(self, file_name: str = "file", content_type: str = "application/octet-stream"):
+        """
+        :param file_name: The name of the file. Note it is overwritten if you use from_file/from_starlette.
+        :param content_type: The content type of the file. Note it is overwritten if you use from_file/from_starlette.
+        """
+        self.content_type = content_type
+        self.file_name = file_name  # the name of the file also when specified in bytesio
+        self._content_buffer = io.BytesIO()
+
+    def from_any(self, data):
+        """
+        Load a file from any supported data type. The file is loaded into the memory as bytes.
+        """
+        # it is already converted
+        if isinstance(data, MediaFile):
+            return data
+
+        # conversion factory
+        if type(data) in [io.BufferedReader, io.BytesIO]:
+            self.from_bytesio_or_handle(data)
+        elif isinstance(data, str):
+            if self._is_valid_file_path(data):
+                self.from_file(data)
+            else:
+                self.from_base64(data)
+        elif isinstance(data, bytes):
+            self.from_bytes(data)
+        elif type(data).__name__ == 'ndarray':
+            self.from_np_array(data)
+        elif data.__module__ == 'starlette.datastructures' and type(data).__name__ == 'UploadFile':
+            self.from_starlette_upload_file(data)
+
+        return self
+
+    def from_bytesio_or_handle(self, buffer: Union[io.BytesIO, BinaryIO, io.BufferedReader], copy: bool = True):
+        """
+        Set the content of the file from a BytesIO or a file handle.
+        :params buffer: The buffer to read from.
+        :params copy: If true, the buffer is completely read to bytes and the bytes copied to this file.
+        """
+        self._reset_buffer()
+        if type(buffer) in [io.BytesIO, io.BufferedReader]:
+            buffer.seek(0)
+            if not copy:
+                self._content_buffer = buffer
+                self._file_info()
+            else:
+                self.from_bytes(buffer.read())
+                buffer.seek(0)
+
+        return self
+
+    def from_bytesio(self, buffer: Union[io.BytesIO, BinaryIO], copy: bool = True):
+        return self.from_bytesio_or_handle(buffer=buffer, copy=copy)
+
+    def from_file(self, path_or_handle: Union[str, io.BytesIO, io.BufferedReader]):
+        """
+        Load a file from a file path, file handle or base64 and convert it to BytesIO.
+        """
+        if type(path_or_handle) in [io.BufferedReader, io.BytesIO]:
+            self.from_bytesio_or_handle(path_or_handle)
+        elif isinstance(path_or_handle, str):
+            # read file from path
+            self.file_name = os.path.basename(path_or_handle)
+            self.content_type = mimetypes.guess_type(self.file_name)[0] or "application/octet-stream"
+            with open(path_or_handle, 'rb') as file:
+                self.from_bytesio_or_handle(file)
+
+        return self
+
+    def from_bytes(self, data: bytes):
+        self._reset_buffer()
+        self._content_buffer.write(data)
+        self._content_buffer.seek(0)
+        self._file_info()
+        return self
+
+    def from_starlette_upload_file(self, starlette_upload_file):
+        """
+        Load a file from a starlette upload file.
+        :param starlette_upload_file:
+        :return:
+        """
+        content = starlette_upload_file.file.read()
+
+        self.file_name = starlette_upload_file.filename
+        self.content_type = starlette_upload_file.content_type
+        self.from_bytes(content)
+        return self
+
+    def from_base64(self, base64_str: str):
+        decoded = self._decode_base_64_if_is(base64_str)
+        if decoded is not None:
+            return self.from_bytes(base64.b64decode(base64_str))
+        else:
+            err_str = base64_str if len(base64_str) <= 50 else base64_str[:50] + "..."
+            raise ValueError(f"Decoding from base64 like string {err_str} was not possible. Check your data.")
+
+    @requires_numpy()
+    def from_np_array(self, np_array: np.array):
+        """
+        Convert a numpy array to a file which is saved as bytes b"\x93NUMPY" into the buffer.
+        """
+        self._reset_buffer()
+        np.save(self._content_buffer, np_array)
+        return self
+
+    def from_dict(self, file_result_json: dict):
+        """
+        Load a file from a dictionary.
+        :param d: The dictionary to load from formatted as FileResult.to_json().
+        """
+        self.file_name = file_result_json["file_name"]
+        self.content_type = file_result_json["content_type"]
+        # ToDo: the from_base64 might overwrite name and content type (ImageFile). Check if this always is intended.
+        self.from_base64(file_result_json["content"])
+        return self
+
+    @requires_numpy()
+    def to_np_array(self, shape=None, dtype=np.uint8):
+        """
+        If file was created with from_np_array it will return the numpy array.
+        Else it will try to convert the file to a numpy array (note this is converted bytes representation of the file).
+        :param shape: The shape of the numpy array. If None it will be returned flat.
+        :param dtype: The dtype of the numpy array. If None it will be uint8.
+        """
+        bytes = self.to_bytes()
+        # check if was saved with np.save so bytes contains NUMPY
+        if bytes.startswith(b"\x93NUMPY"):
+            self._content_buffer.seek(0)
+            return np.load(self._content_buffer, allow_pickle=False)
+
+        shape = shape or (1, len(bytes))
+        dtype = dtype or np.uint8
+
+        arr_flat = np.frombuffer(bytes, dtype=dtype)
+        return arr_flat.reshape(shape)
+
+    def to_bytes(self) -> bytes:
+        return self.read()
+
+    def read(self) -> bytes:
+        self._content_buffer.seek(0)
+        res = self._content_buffer.read()
+        self._content_buffer.seek(0)
+        return res
+
+    def to_bytes_io(self) -> io.BytesIO:
+        return self._content_buffer
+
+    def to_base64(self):
+        return base64.b64encode(self.to_bytes()).decode()
+
+    def to_httpx_send_able_tuple(self):
+        return self.file_name, self.read(), self.content_type
+
+    def _reset_buffer(self):
+        self._content_buffer.seek(0)
+        self._content_buffer.truncate(0)
+
+    def save(self, path: str):
+        """
+        Methods saves the file to disk.
+        If path is a folder it will save it in folder/self.filename.
+        If path is a file it will save it there.
+        :param path:
+        :return:
+        """
+        # create folder if not exists
+        if os.path.dirname(path) != "" and not os.path.exists(os.path.dirname(path)):
+            os.makedirs(os.path.dirname(path))
+        # check if path contains a file name
+        if os.path.basename(path) == "":
+            path = os.path.join(path, self.file_name)
+
+        with open(path, 'wb') as file:
+            file.write(self.read())
+
+    def _file_info(self):
+        """
+        After writing the file to the buffer, this method is called to determine additional file informations.
+        For videos this might be length, frame rate...
+        If you subclass don't forget to call super()._file_info() to set the file name and content type.
+        """
+        # set file name and type
+        if hasattr(self._content_buffer, "name"):
+            self.file_name = os.path.basename(self._content_buffer.name)
+
+        if self.file_name != "file":
+            self.content_type = mimetypes.guess_type(self.file_name)[0] or "application/octet-stream"
+        else:
+            self.content_type = "application/octet-stream"
+
+    def __bytes__(self):
+        return self.to_bytes()
+
+    def __array__(self):
+        return self.to_np_array()
+
+    def to_json(self):
+        """
+        Returns the file as a json serializable dictionary.
+        :return: { "file_name": str, "content_type": str, "content": str }
+        """
+        return {
+            "file_name": self.file_name,
+            "content_type": self.content_type,
+            "content": self.to_base64()
+        }
+
+    @staticmethod
+    def _decode_base_64_if_is(data: Union[bytes, str]):
+        """
+        Checks if a string is base64. If it is, it returns the base64 string as bytes; else returns None.
+        """
+        if isinstance(data, str):
+            data = data.encode()
+
+        try:
+            # Decode the data
+            decoded = base64.b64decode(data, validate=True)
+            # Re-encode the decoded data
+            back_encoded = base64.b64encode(decoded)
+            # Compare with the original encoded data
+            if back_encoded == data:
+                return decoded
+        except Exception:
+            pass
+
+        return None
+
+    @staticmethod
+    def _is_valid_file_path(path: str):
+        try:
+            is_file = os.path.isfile(path)
+            return is_file
+        except:
+            return False
+

media_toolkit/core/video/video_file.py

@@ -0,0 +1,283 @@
+import glob
+import os
+import sys
+import tempfile
+from io import BytesIO
+from typing import List, Union
+
+from media_toolkit.core.video.video_utils import (add_audio_to_video_file, audio_array_to_audio_file,
+                                                     video_from_image_generator, get_sample_rate_from_audio_file)
+from media_toolkit.dependency_requirements import requires
+from media_toolkit.core.media_file import MediaFile
+
+
+try:
+    import cv2
+    import numpy as np
+except ImportError:
+    pass
+
+try:
+    from vidgear.gears import VideoGear, WriteGear
+except:
+    pass
+
+try:
+    from pydub import AudioSegment
+    from pydub.utils import mediainfo
+except ImportError:
+    pass
+
+
+class VideoFile(MediaFile):
+    """
+    A class to represent a video file.
+    """
+    def __init__(self):
+        super().__init__()
+        self.content_type = "video"
+        self.frame_count = None  # an estimated value based on cv2.VideoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
+        self.frame_rate = None
+        self.shape = None
+        self.audio_sample_rate = None
+
+    def from_files(self, image_files: Union[List[str], list], frame_rate: int = 30, audio_file=None):
+        """
+        Creates a video based of a list of image files and an audio_file file.
+        :param image_files: A list of image files to convert to a video. Either paths or numpy arrays.
+        :param frame_rate: The frame rate of the video.
+        :param audio_file: The audio_file file to add to the video, as path, bytes or AudioSegment.
+        """
+        # Check if there are images in the list
+        if not image_files:
+            raise ValueError("The list of image files is empty.")
+
+        # Create a temporary file to store the video
+        temp_vid_file_path = video_from_image_generator(image_files, frame_rate=frame_rate, save_path=None)
+        # Merge video and audio_file using pydub
+        if audio_file is not None:
+            combined = add_audio_to_video_file(video_file=temp_vid_file_path, audio_file=audio_file)
+            self.from_file(combined)
+            os.remove(combined)
+            os.remove(temp_vid_file_path)
+            return self
+
+        # Init self from the temp file
+        self.from_file(temp_vid_file_path)
+        # remove tempfile
+        os.remove(temp_vid_file_path)
+
+        return self
+
+    def from_image_files(self, image_files: List[str], frame_rate: int = 30):
+        """
+        Converts a list of image files into a video file.
+        """
+        return self.from_files(image_files, frame_rate, audio_file=None)
+
+    def from_dir(self, dir_path: str, audio: Union[str, list] = None, frame_rate: int = 30):
+        """
+        Converts all images in a directory into a video file.
+        """
+        image_types = ["*.png", "*.jpg", "*.jpeg"]
+        image_files = []
+        for image_type in image_types:
+            image_files.extend(glob.glob(os.path.join(dir_path, image_type)))
+        # sort by date to make sure the order is correct
+        image_files.sort(key=lambda x: os.path.getmtime(x))
+
+        # if audio_file is none, take the first audio_file file in the directory
+        if audio is None:
+            audio_types = ["*.wav", "*.mp3"]
+            for audio_type in audio_types:
+                audio = glob.glob(os.path.join(dir_path, audio_type))
+                if len(audio) > 0:
+                    audio = audio[0]
+                    break
+
+        return self.from_files(image_files=image_files, frame_rate=frame_rate, audio_file=audio)
+
+    def add_audio(self, audio_file: Union[str, list], sample_rate: int = 44100):
+        if self.audio_sample_rate is None:
+            if self.frame_rate is None:
+                raise Exception("The frame rate of the video file is not set. Read a video file first.")
+
+            if os.path.isfile(audio_file):
+                self.audio_sample_rate = get_sample_rate_from_audio_file(audio_file)
+            else:
+
+                self.audio_sample_rate = int(mediainfo(self._to_temp_file())['sample_rate'])
+
+
+        if isinstance(audio_file, list) or isinstance(audio_file, np.ndarray):
+            audio_file = audio_array_to_audio_file(audio_file, sample_rate=self.audio_sample_rate)
+
+        tmp = self._to_temp_file()
+        combined = add_audio_to_video_file(tmp, audio_file)
+        self.from_file(combined)
+        os.remove(tmp)
+        os.remove(combined)
+        return self
+
+    def _to_temp_file(self):
+        # get suffix
+        if self.content_type is None:
+            raise ValueError("The content type of the video file is not set.")
+        if "/" not in self.content_type:
+            raise ValueError("The content type of the video file is not valid. Read a video file first.")
+        suffix = self.content_type.split("/")[1]
+        if suffix == 'octet-stream':
+            raise ValueError("The content type of the video file is not valid. Read a video file first.")
+
+        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{suffix}") as temp_video_file:
+            temp_video_file.write(self._content_buffer.getvalue())
+            temp_video_file_path = temp_video_file.name
+
+        return temp_video_file_path
+
+    @requires('vidgear', 'numpy', 'pydub')
+    def from_video_stream(self, video_audio_stream, frame_rate: int = 30):
+        """
+        Given a generator that yields video frames and audio_file data as numpy arrays, this creates a video.
+        The generator is expected to be in the form of: VideoFile().to_video_stream()
+        """
+        # Reset and pre-settings
+        self._reset_buffer()
+
+        # new generator, to extract audio_file
+        audio_frames = []
+        def _frame_gen():
+            for frame in video_audio_stream:
+                # check if is video and audio_file stream or only video stream
+                if len(frame) == 2:
+                    frame, audio_data = frame
+                    audio_frames.append(audio_data)
+            yield frame
+
+        # Create video
+        temp_video_file_path = video_from_image_generator(_frame_gen(), frame_rate=frame_rate, save_path=None)
+
+        # Add audio_file
+        if len(audio_frames) > 0:
+            temp_audio_file = audio_array_to_audio_file(audio_frames, sample_rate=self.audio_sample_rate)
+            combined = add_audio_to_video_file(temp_video_file_path, temp_audio_file)
+            self.from_file(combined)
+            # cleanup
+            os.remove(temp_audio_file)
+            os.remove(temp_video_file_path)
+            os.remove(combined)
+
+    @requires('cv2', 'pydub')
+    def _file_info(self):
+        super()._file_info()
+
+        #if file_path is not None:
+        #    temp = file_path
+        #else:
+        self._content_buffer.seek(0)
+        temp = self._to_temp_file()
+
+        cap = cv2.VideoCapture(temp)
+        self.frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # is an estimated value.
+        # determine content codec
+        # https://stackoverflow.com/questions/61659346/how-to-get-4-character-codec-code-for-videocapture-object-in-opencv
+        # h = int(cap.get(cv2.CAP_PROP_FOURCC))
+        # b = h.to_bytes(4, byteorder=sys.byteorder)
+        # codec = b.decode()  # results in the codec
+        self.content_type = f"video/mp4"
+        self.frame_rate = cap.get(cv2.CAP_PROP_FPS)
+        self.shape = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+        cap.release()
+        # get audio sample rate
+        info = mediainfo(temp)
+        if 'sample_rate' in info:
+            try:
+                self.audio_sample_rate = int(info['sample_rate'])
+            except ValueError:
+                self.audio_sample_rate = 44100
+        os.remove(temp)
+
+    @requires('vidgear')
+    def to_image_stream(self):
+        return self.to_video_stream(include_audio=False)
+
+    @requires('pydub', 'vidgear')
+    def to_video_stream(self, include_audio=True):
+        """
+        Yields video frames and audio_file data as numpy arrays.
+        :param include_audio: if the audio_file is included in the video stream. If not it will only yield the video frames.
+        :return:
+        """
+        self._content_buffer.seek(0)
+        # because CamGear does not support reading from a BytesIO buffer, we need to save the buffer to a temporary file
+        temp_video_file_path = self._to_temp_file()
+        stream = VideoGear(source=temp_video_file_path).start()
+
+        if include_audio:
+            # Extract audio_file using pydub
+            audio = AudioSegment.from_file(temp_video_file_path)
+            # Calculate the audio_file segment duration per frame
+            audio_per_frame_duration = 1000 / stream.framerate  # duration of each video frame in ms
+            # Initialize frame counter for audio_file
+            frame_count = 0
+
+        try:
+            while True:
+                # Read frame
+                frame = stream.read()
+                if frame is None:
+                    break
+
+                if not include_audio:
+                    yield frame
+                    continue
+
+                # Calculate the start and end times for the corresponding audio_file segment
+                start_time = frame_count * audio_per_frame_duration
+                end_time = start_time + audio_per_frame_duration
+                frame_audio = audio[start_time:end_time]
+
+                # Convert audio_file segment to raw data
+                audio_data = np.array(frame_audio.get_array_of_samples())
+
+                # Yield the frame and the corresponding audio_file data
+                yield frame, audio_data
+
+                # Increment frame counter
+                frame_count += 1
+        finally:
+            # Safely close the video stream
+            stream.stop()
+            # Remove the temporary video file
+            os.remove(temp_video_file_path)
+            # accurate value instead of using cv2.CAP_PROP_FRAME_COUNT
+            self.frame_count = frame_count
+
+    @requires('pydub')
+    def extract_audio(self, path: str = None, export_type: str = 'mp4') -> Union[bytes, None]:
+        temp_video_file_path = self._to_temp_file()
+        audio = AudioSegment.from_file(temp_video_file_path)
+
+        if path is not None and len(path) > 0:
+            dirname = os.path.dirname(path)
+            if len(dirname) > 0 and not os.path.isdir(dirname):
+                os.makedirs(dirname)
+            audio.export(path, format=export_type)
+            os.remove(temp_video_file_path)
+            return None
+
+        # return as bytes
+        file = BytesIO()
+        file = audio.export(file, format=export_type)
+        file.seek(0)
+        data = file.read()
+        file.close()
+        # remove tempfile
+        os.remove(temp_video_file_path)
+        return data
+
+    def __iter__(self):
+        return self.to_video_stream()
+
+    def __len__(self):
+        return self.frame_count