Add initial support for USM alpha chunks

Author: donmai-me

CHANGELOG.md

@@ -4,9 +4,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## Unreleased
+## [0.2.5] - 2022-04-08
 ### Added
 - New operation in command-line called `encryptusm` which encrypts an existing USM file.
+- Support for new USM Chunk `@ALP` which is used for alpha transparency videos. Currently, supports probe and extraction operations only. Thank you to [EmirioBomb](https://github.com/EmirioBomb) for bringing this to my attention and providing sample files.
+
+### Changed
+- Renamed Usm constructor parameters.
 
 ### Removed
 - Check for chunk header in `chunk_size_and_padding` function
@@ -49,7 +53,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed bug in extractusm that causes it to fail when output directory doesn't exist.
 - Fixed bug where program fails when directory exists.
 
-[Unreleased]: https://github.com/donmai-me/WannaCRI/compare/0.2.4...HEAD
+[Unreleased]: https://github.com/donmai-me/WannaCRI/compare/0.2.5...HEAD
+[0.2.5]: https://github.com/donmai-me/WannaCRI/compare/0.2.4...0.2.5
 [0.2.4]: https://github.com/donmai-me/WannaCRI/compare/0.2.3...0.2.4
 [0.2.3]: https://github.com/donmai-me/WannaCRI/compare/0.2.2...0.2.3
 [0.2.2]: https://github.com/donmai-me/WannaCRI/compare/0.2.1...0.2.2

wannacri/usm/media/protocols.py

@@ -52,10 +52,21 @@ def metadata_pages(self, pages: Optional[List[UsmPage]]):
     def channel_number(self) -> int:
         return self._channel_number
 
+    @channel_number.setter
+    def channel_number(self, new_channel_number: int):
+        if new_channel_number < 0:
+            raise ValueError(f"Given negative channel number: {new_channel_number}")
+
+        self._channel_number = new_channel_number
+
     @property
     def header_page(self) -> UsmPage:
         return self._header_page
 
+    @header_page.setter
+    def header_page(self, new_header_page: UsmPage):
+        self._header_page = new_header_page
+
     @property
     def filename(self) -> str:
         """A slugified filename, which is stored inside the crid_page.
@@ -68,6 +79,12 @@ def filename(self) -> str:
         result: str = filename.val.split("/")[-1]
         return slugify(result, allow_unicode=True)
 
+    @filename.setter
+    def filename(self, new_filename: str):
+        new_filename = slugify(new_filename, allow_unicode=True)
+        self.crid_page["filename"].type = ElementType.STRING
+        self.crid_page["filename"].val = new_filename
+
     def __len__(self) -> int:
         """The number of packets a Usm videos or audios has."""
         return self._length
@@ -89,6 +106,7 @@ class UsmVideo(UsmMedia, Protocol):
     # Classes that explicitly inherit UsmVideo should have this attribute
     # to use the default stream and chunks methods.
     _stream: Generator[Tuple[bytes, bool], None, None]
+    is_alpha: bool
 
     def stream(
         self, mode: OpMode = OpMode.NONE, key: Optional[bytes] = None
@@ -317,246 +335,3 @@ def chunks(
                         channel_number=self.channel_number,
                     ),
                 ]
-
-
-# TODO: Delete the comments when done rewriting
-
-# class VP9:
-#     def __init__(self, filepath: str):
-#         self.filename = os.path.basename(filepath)
-#         self.filesize = os.path.getsize(filepath)
-#         self.info = ffmpeg.probe(filepath, show_entries="packet=dts,pts_time,pos,flags")
-
-#         if len(self.info.get("streams")) == 0:
-#             raise Exception("No streams found")
-#         elif len(self.info.get("streams")) > 1:
-#             raise NotImplementedError("Can only accept one stream")
-#         elif self.info.get("format").get("format_name") != "ivf":
-#             raise Exception("Not an ivf file")
-#         elif self.info.get("streams")[0].get("codec_name") != "vp9":
-#             raise Exception("Not a VP9 videos")
-
-#         video_stream = self.info.get("streams")[0]
-#         self.bitrate = int(self.info.get("format").get("bit_rate"))
-#         self.width = int(video_stream.get("width"))
-#         self.height = int(video_stream.get("height"))
-#         self.framerate_n = int(video_stream.get("r_frame_rate").split("/")[0])
-#         self.framerate_d = int(video_stream.get("r_frame_rate").split("/")[1])
-
-#         self.frames = self.info.get("packets")
-#         self.keyframes = [frame for frame in self.frames if frame.get("flags") == "K_"]
-#         self.total_frames = len(self.frames)
-#         self.file = open(filepath, "rb")
-
-#     def export(
-#         self, encrypt: bool, key: Optional[int] = None, encoding: str = "UTF-8"
-#     ) -> bytes:
-#         if encrypt:
-#             video_key1, video_key2, _ = generate_keys(key)
-
-#         debug = open("FRAME_TIME_DEBUG", "w+")
-#         stream_chunks = bytearray()
-#         keyframe_usm_offsets = []
-#         max_frame_size = 0
-#         max_packed_frame_size = 0
-#         max_keyframe_to_keyframe_size = 0
-#         current_keyframe_to_keyframe_size = 0
-#         self.file.seek(0, 0)
-#         for i, frame in enumerate(self.frames):
-#             # frame_time formula is based on existing usm files.
-#             # TODO: Does this hold up for videos that's not 30fps?
-#             debug.write("Frame {}: time = {}".format(i, frame.get("pts_time")))
-#             # frame_time = int(99.86891 * i)
-#             frame_time = int(i * 99.9)
-#             if frame in self.keyframes:
-#                 keyframe_usm_offsets.append(len(stream_chunks))
-
-#             if i == len(self.frames) - 1:
-#                 # Last frame
-#                 frame_size = self.filesize - int(frame.get("pos"))
-#             else:
-#                 frame_size = int(self.frames[i + 1].get("pos")) - int(frame.get("pos"))
-
-#             if i == 0:
-#                 # Include 32 byte header for first frame
-#                 max_frame_size = frame_size + 32
-#                 packet = self.file.read(frame_size + 32)
-#             else:
-#                 if frame_size > max_frame_size:
-#                     max_frame_size = frame_size
-
-#                 packet = self.file.read(frame_size)
-
-#             if frame.get("flags") != "K_":
-#                 current_keyframe_to_keyframe_size += frame_size
-#             elif current_keyframe_to_keyframe_size > max_keyframe_to_keyframe_size:
-#                 max_keyframe_to_keyframe_size = current_keyframe_to_keyframe_size
-#                 current_keyframe_to_keyframe_size = frame_size
-
-#             if encrypt:
-#                 packet = encrypt_video_packet(packet, video_key1, video_key2)
-
-#             padding_size = 0x20 - (len(packet) % 0x20) if len(packet) % 0x20 != 0 else 0
-#             packed_frame_chunk = generate_packed_chunk(
-#                 "@SFV",
-#                 0,
-#                 int(100 * self.framerate_n / self.framerate_d),
-#                 frame_time,
-#                 packet,
-#                 padding_size,
-#             )
-#             if len(packed_frame_chunk) > max_packed_frame_size:
-#                 max_packed_frame_size = len(packed_frame_chunk)
-
-#             stream_chunks += packed_frame_chunk
-
-#         stream_chunks += generate_packed_chunk(
-#             "@SFV",
-#             2,
-#             int(self.framerate_n / self.framerate_d),
-#             0,
-#             bytes("#CONTENTS END   ===============", "UTF-8") + bytes(1),
-#             0,
-#         )
-
-#         seek_chunks = bytearray()
-#         keyframe_pages = []
-#         video_header_end_offset = get_video_header_end_offset(len(self.keyframes))
-#         # Offset of videos header end offset plus length of videos header end
-#         stream_offset = video_header_end_offset + 0x40
-#         for i, keyframe in enumerate(self.keyframes):
-#             keyframe_usm_offset = keyframe_usm_offsets[i]
-#             keyframe_offset = stream_offset + keyframe_usm_offset
-#             keyframe_page = UsmPage("VIDEO_SEEKINFO", i)
-#             keyframe_page.add("ofs_byte", ElementType.clonglong, keyframe_offset)
-#             keyframe_page.add("ofs_frmid", ElementType.cuint, keyframe.get("dts"))
-#             keyframe_page.add("num_skip", ElementType.cushort, 0)
-#             keyframe_page.add("resv", ElementType.cushort, 0)
-#             keyframe_pages.append(keyframe_page)
-
-#         seek_chunks_payload = pack_pages(keyframe_pages, string_padding=1)
-#         # 0x20 bytes for chunk header.
-#         seek_chunks_padding = (
-#             video_header_end_offset - 0xA40 - 0x20 - len(seek_chunks_payload)
-#         )
-#         seek_chunks += generate_packed_chunk(
-#             "@SFV",
-#             3,
-#             int(self.framerate_n / self.framerate_d),
-#             0,
-#             seek_chunks_payload,
-#             seek_chunks_padding,
-#         )
-#         metadata_size = len(seek_chunks)
-#         seek_chunks += generate_packed_chunk(
-#             "@SFV",
-#             2,
-#             int(self.framerate_n / self.framerate_d),
-#             0,
-#             bytes("#METADATA END   ===============", "UTF-8") + bytes(1),
-#             0,
-#         )
-
-#         header_page = UsmPage("VIDEO_HDRINFO", 0)
-#         header_page.add("width", ElementType.cint, self.width)
-#         header_page.add("height", ElementType.cint, self.height)
-#         header_page.add("mat_width", ElementType.cint, self.width)
-#         header_page.add("mat_height", ElementType.cint, self.height)
-#         header_page.add("disp_width", ElementType.cint, self.width)
-#         header_page.add("disp_height", ElementType.cint, self.height)
-#         header_page.add("scrn_width", ElementType.cint, 0)
-#         header_page.add("mpeg_dcprec", ElementType.cchar, 0)
-#         header_page.add("mpeg_codec", ElementType.cchar, 9)
-#         # TODO: Check if videos has transparency
-#         header_page.add("alpha_type", ElementType.cint, 0)
-#         header_page.add("total_frames", ElementType.cint, self.total_frames)
-
-#         framerate_n = self.framerate_n
-#         framerate_d = self.framerate_d
-
-#         if framerate_d < 1000 and framerate_d != 1000:
-#             framerate_d *= 1000
-#             framerate_n *= 1000
-
-#         header_page.add("framerate_n", ElementType.cint, framerate_n)
-#         header_page.add("framerate_d", ElementType.cint, framerate_d)
-#         header_page.add("metadata_count", ElementType.cint, 1)
-#         header_page.add("metadata_size", ElementType.cint, metadata_size)
-#         header_page.add("ixsize", ElementType.cint, max_packed_frame_size)
-#         header_page.add("pre_padding", ElementType.cint, 0)
-#         header_page.add("max_picture_size", ElementType.cint, 0)
-#         header_page.add("color_space", ElementType.cint, 0)
-#         header_page.add("picture_type", ElementType.cint, 0)
-
-#         header_chunk_payload = pack_pages([header_page])
-#         header_chunk_padding = 0xA00 - 0x800 - 0x20 - len(header_chunk_payload)
-#         header_chunk = bytearray()
-#         header_chunk += generate_packed_chunk(
-#             "@SFV",
-#             1,
-#             int(self.framerate_n / self.framerate_d),
-#             0,
-#             header_chunk_payload,
-#             header_chunk_padding,
-#         )
-#         header_chunk += generate_packed_chunk(
-#             "@SFV",
-#             2,
-#             int(self.framerate_n / self.framerate_d),
-#             0,
-#             bytes("#HEADER END     ===============", "UTF-8") + bytes(1),
-#             0,
-#         )
-
-#         directory_pages = []
-#         for i in range(0, 2):
-#             if i == 0:
-#                 filename = os.path.splitext(self.filename)[0] + ".usm"
-#                 # filename = r"I:\000125 千本桜\000125.usm"
-#                 filesize = (
-#                     0x800 + len(header_chunk) + len(seek_chunks) + len(stream_chunks)
-#                 )
-#                 stmid = 0
-#                 chno = -1
-#                 minchk = 1
-#                 # TODO: Find formula for minbuf for usm
-#                 minbuf = round(1.98746 * max_frame_size)
-#                 minbuf += 0x10 - (minbuf % 0x10) if minbuf % 0x10 != 0 else 0
-#             else:
-#                 filename = self.filename
-#                 # filename = r"I:\000125 千本桜\000125.ivf"
-
-#                 filesize = self.filesize
-#                 stmid = 1079199318  # @SFV
-#                 chno = 0
-#                 minchk = 3
-#                 minbuf = max_frame_size
-
-#             directory_part = UsmPage("CRIUSF_DIR_STREAM", 0)
-#             directory_part.add("fmtver", ElementType.cint, 16777984)
-#             directory_part.add("filename", ElementType.cstring, filename)
-#             directory_part.add("filesize", ElementType.cint, filesize)
-#             directory_part.add("datasize", ElementType.cint, 0)
-#             directory_part.add("stmid", ElementType.cint, stmid)
-#             directory_part.add("chno", ElementType.cshort, chno)
-#             directory_part.add("minchk", ElementType.cshort, minchk)
-#             directory_part.add("minbuf", ElementType.cint, minbuf)
-#             directory_part.add("avbps", ElementType.cint, self.bitrate)
-#             directory_pages.append(directory_part)
-
-#         directory_chunk_payload = pack_pages(directory_pages, encoding, 5)
-#         directory_chunk_padding = 0x800 - 0x20 - len(directory_chunk_payload)
-#         directory_chunk = bytearray()
-#         directory_chunk += generate_packed_chunk(
-#             "CRID",
-#             1,
-#             int(self.framerate_n / self.framerate_d),
-#             0,
-#             directory_chunk_payload,
-#             directory_chunk_padding,
-#         )
-#         result = directory_chunk + header_chunk + seek_chunks + stream_chunks
-#         return bytes(result)
-
-#     def close(self):
-#         self.file.close()

wannacri/usm/media/video.py

@@ -21,13 +21,15 @@ def __init__(
         length: int,
         channel_number: int = 0,
         metadata_pages: Optional[List[UsmPage]] = None,
+        is_alpha: bool = False,
     ):
         self._stream = stream
         self._crid_page = crid_page
         self._header_page = header_page
         self._length = length
         self._channel_number = channel_number
         self._metadata_pages = metadata_pages
+        self.is_alpha = is_alpha
 
 
 class Vp9(UsmVideo):

wannacri/usm/tools.py

@@ -39,6 +39,7 @@ def is_valid_chunk(signature: bytes) -> bool:
         bytes("CRID", "UTF-8"),  # CRI USF DIR STREAM
         bytes("@SFV", "UTF-8"),  # Video
         bytes("@SFA", "UTF-8"),  # Audio
+        bytes("@ALP", "UTF-8")
     ]
     return signature[:4] in valid_signatures
 

wannacri/usm/types.py

@@ -8,6 +8,7 @@ class ChunkType(Enum):
     INFO = bytearray("CRID", "UTF-8")
     VIDEO = bytearray("@SFV", "UTF-8")
     AUDIO = bytearray("@SFA", "UTF-8")
+    ALPHA = bytearray("@ALP", "UTF-8")
 
     @staticmethod
     def from_bytes(data: bytes) -> ChunkType:
@@ -19,6 +20,9 @@ def from_bytes(data: bytes) -> ChunkType:
 
         raise ValueError(f"Unknown chunk signature: {bytes_to_hex(data[:4])}")
 
+    def __str__(self):
+        return str(self.value, "UTF-8")
+
 
 class PayloadType(Enum):
     STREAM = 0
@@ -61,7 +65,7 @@ class ElementType(Enum):
     ULONGLONG = 0x17  # 8 bytes
     FLOAT = 0x18  # 4 bytes
     # TODO: Confirm DOUBLE's existence
-    # DOUBLE = 0x19 # 8 bytes
+    # DOUBLE = 0x19  # 8 bytes
     STRING = 0x1A  # Null byte terminated
     BYTES = 0x1B  # Bytes