mirror of
https://github.com/Akascape/Datamosher-Pro.git
synced 2025-12-05 15:59:59 +01:00
125 lines
3.5 KiB
Python
125 lines
3.5 KiB
Python
import struct
|
|
|
|
from pymosh.codec.mpeg4 import is_iframe
|
|
|
|
from . import riff
|
|
|
|
|
|
class AVIFile(object):
|
|
"""A wrapper for AVI files."""
|
|
|
|
def __init__(self):
|
|
self.riff = riff.RiffIndex()
|
|
self.streams = []
|
|
self.frame_order = []
|
|
|
|
@staticmethod
|
|
def from_file(filename: str):
|
|
instance = AVIFile()
|
|
|
|
instance.riff = riff.RiffIndex.from_file(filename=filename)
|
|
|
|
header = instance.riff.find(b'LIST', b'hdrl')
|
|
# Get stream info
|
|
stream_lists = header.find_all(b'LIST', b'strl')
|
|
for l in stream_lists:
|
|
strh = l.find(b'strh')
|
|
data = strh.data
|
|
fccType, = struct.unpack(b'4s', data[:4])
|
|
stream = Stream(len(instance.streams), fccType)
|
|
instance.streams.append(stream)
|
|
|
|
instance.split_streams()
|
|
|
|
return instance
|
|
|
|
def __iter__(self):
|
|
return iter(self.streams)
|
|
|
|
def add_frame(self, chunk):
|
|
stream_num = int(chunk.header[:2])
|
|
if stream_num < len(self.streams):
|
|
self.frame_order.append(
|
|
(stream_num, len(self.streams[stream_num])))
|
|
self.streams[stream_num].add_frame(chunk)
|
|
|
|
def split_streams(self):
|
|
movi = self.riff.find(b'LIST', b'movi')
|
|
for chunk in movi:
|
|
self.add_frame(chunk)
|
|
|
|
def combine_streams(self):
|
|
chunks = []
|
|
for frame_record in self.frame_order:
|
|
stream_num, frame_num = frame_record
|
|
stream = self.streams[stream_num]
|
|
frame = stream[frame_num]
|
|
chunks.append(frame)
|
|
return chunks
|
|
|
|
def _video(self):
|
|
return filter(lambda stream: stream.type == b'vids', self.streams)
|
|
video = property(_video)
|
|
|
|
def _audio(self):
|
|
return filter(lambda stream: stream.type == b'auds', self.streams)
|
|
audio = property(_audio)
|
|
|
|
def rebuild(self):
|
|
"""Rebuild RIFF tree and index from streams."""
|
|
movi = self.riff.find(b'LIST', b'movi')
|
|
movi.chunks = self.combine_streams()
|
|
self.rebuild_index()
|
|
|
|
def rebuild_index(self):
|
|
old_index = self.riff.find(b'idx1')
|
|
movi = self.riff.find(b'LIST', b'movi')
|
|
data = b''
|
|
offset = 0
|
|
flags = {
|
|
'base': 0x00000000,
|
|
'keyframe': 0x00000010,
|
|
}
|
|
for chunk in movi:
|
|
length = len(chunk)
|
|
frame_flags = flags['base']
|
|
# If it's a video keyframe or audio frame, use keyframe flag
|
|
if (chunk.header[2] == b'd' and is_iframe(chunk)) or (chunk.header[2] == b'w'):
|
|
frame_flags |= flags['keyframe']
|
|
data += struct.pack(b'<4sIII', chunk.header, frame_flags, offset,
|
|
length+8)
|
|
offset += length + 8 + (length % 2)
|
|
new_index = riff.RiffDataChunk(b'idx1', data)
|
|
self.riff.find(b'RIFF').replace(old_index, new_index)
|
|
|
|
def write(self, fh):
|
|
self.riff.write(fh)
|
|
|
|
|
|
class Stream(object):
|
|
def __init__(self, num, stream_type):
|
|
self.num = int(num)
|
|
self.type = stream_type
|
|
self.chunks = []
|
|
|
|
def add_frame(self, chunk):
|
|
self.chunks.append(chunk)
|
|
|
|
def __getitem__(self, index):
|
|
return self.chunks.__getitem__(index)
|
|
|
|
def __iter__(self):
|
|
return self.chunks.__iter__()
|
|
|
|
def __len__(self):
|
|
return len(self.chunks)
|
|
|
|
def append(self, *args):
|
|
return self.chunks.append(*args)
|
|
|
|
def extend(self, *args):
|
|
return self.chunks.extend(*args)
|
|
|
|
def replace(self, chunks):
|
|
self.chunks = chunks
|