Commit cc192210 authored by norvell's avatar norvell
Browse files

Update cut_bs.py such that --frame and --length may be combined

parent e8f8c027
Loading
Loading
Loading
Loading
Loading
+36 −28
Original line number Diff line number Diff line
@@ -59,45 +59,42 @@ def cut_to_length(fp, fp_out, length):

    return fr_cnt


def cut_from_start(fp, fp_out, start_frame=0, start_with_sid=False):
def cut_from_start(fp, fp_out, start_frame=0, length=-1, start_with_sid=False):
    # cut until start frame
    fr_cnt = 0
    cut_cnt = 0
    extracted_frames_cnt = 0
    for cur_frame in range(start_frame):
        sync_word = fp.read(2)
        if sync_word == b"":
            return (fr_cnt, cut_cnt)
            return (fr_cnt, extracted_frames_cnt)
        if sync_word not in SYNC_WORDS:
            raise ValueError("Bad Sync word!")
        n_bits = struct.unpack("h", fp.read(2))[0]
        fp.read(n_bits * 2)
        fr_cnt += 1
        cut_cnt += 1
    if start_with_sid:
        found = False
        while not found:
            sync_word = fp.read(2)
            if sync_word == b"":
                return (fr_cnt, cut_cnt)
                return (fr_cnt, extracted_frames_cnt)
            if sync_word not in SYNC_WORDS:
                raise ValueError("Bad Sync word!")
            n_bits_bs = fp.read(2)
            n_bits = struct.unpack("h", n_bits_bs)[0]
            if n_bits in SID_BITS:
                found = True
                fp_out.write(sync_word)
                fp_out.write(n_bits_bs)
                fp_out.write(fp.read(n_bits * 2))
                fp.seek(-4,1) # Found SID frame, go back
            else:
                fp.read(n_bits * 2)
                cut_cnt += 1
                fr_cnt += 1

    # transfer remaining frames
    while True:
    # transfer frames
    eof = False
    while length != 0:
        sync_word = fp.read(2)
        if sync_word == b"":
            eof = True
            break
        if sync_word not in SYNC_WORDS:
            raise ValueError("Bad Sync word!")
@@ -107,7 +104,23 @@ def cut_from_start(fp, fp_out, start_frame=0, start_with_sid=False):
        fp_out.write(n_bits_bs)
        fp_out.write(fp.read(n_bits * 2))
        fr_cnt += 1
    return (fr_cnt, cut_cnt)
        extracted_frames_cnt += 1
        length -= 1

    # count remaining frames, if any
    while not eof:
        sync_word = fp.read(2)
        if sync_word == b"":
            eof = True
            break
        if sync_word not in SYNC_WORDS:
            raise ValueError("Bad Sync word!")
        n_bits_bs = fp.read(2)
        n_bits = struct.unpack("h", n_bits_bs)[0]
        fp.read(n_bits * 2)
        fr_cnt += 1

    return (fr_cnt, extracted_frames_cnt)


if __name__ == "__main__":
@@ -145,17 +158,12 @@ if __name__ == "__main__":

    with open(my_args.bs_in, "rb") as fp_in:
        with open(my_args.bs_out, "wb") as fp_out:
            if my_args.sid or my_args.frame:
                fr_cnt, cut_cnt = cut_from_start(
                    fp_in, fp_out, start_frame=my_args.frame, start_with_sid=my_args.sid
            fr_cnt, extracted_frames_cnt = cut_from_start(
                fp_in, fp_out, start_frame=my_args.frame, length=my_args.length, start_with_sid=my_args.sid
            )
                if my_args.sid and (fr_cnt == cut_cnt):
                    print("Warning! No SID frame found in bitstream!")
                print(f"Cut {cut_cnt} of {fr_cnt} frames from {my_args.bs_in}")
            elif my_args.length:
                fr_cnt = cut_to_length(fp_in, fp_out, my_args.length)
                if fr_cnt != my_args.length:
                    print(
                        f"Warning! Could not cut to length {my_args.length} as bitstream only contained {fr_cnt} frames!"
                    )
                print(f"Cut {my_args.bs_in} to {fr_cnt} frames")
            if extracted_frames_cnt == 0:
                print(f"\nExtracted 0 frames!")
                sys.exit(-1)
            if (my_args.length != -1 and extracted_frames_cnt != my_args.length):
                print(f"Warning! Requested {my_args.length}, only {extracted_frames_cnt} obtained!")
            print(f"Extracted {extracted_frames_cnt} out of {fr_cnt} frames from {my_args.bs_in}")