Source file speex.ml

(*
 * Copyright 2003-2011 Savonet team
 *
 * This file is part of Ocaml-speex.
 *
 * Ocaml-speex is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * Ocaml-speex is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Ocaml-speex; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

(**
  * Functions for decoding speex files using libspeex.
  *
  * @author Romain Beauxis
  *)

exception Invalid_frame_size

type mode = Narrowband | Wideband | Ultra_wideband

(* Internal use only *)
type internal_mode

external internal_mode_of_int : int -> internal_mode = "caml_speex_get_mode"

(* Values from speex.h *)
let mode_of_int x =
  match x with
    | 0 -> Narrowband
    | 1 -> Wideband
    | 2 -> Ultra_wideband
    | _ -> failwith "unknown mode"

let int_of_mode x =
  match x with Narrowband -> 0 | Wideband -> 1 | Ultra_wideband -> 2

let internal_mode_of_mode x = internal_mode_of_int (int_of_mode x)

(* Generated by control_define *)

(** TODO: better implementation.. *)
type control =
  | SPEEX_RESET_STATE
  | SPEEX_SET_ENH
  | SPEEX_GET_ENH
  | SPEEX_GET_FRAME_SIZE
  | SPEEX_SET_QUALITY
  | SPEEX_SET_MODE
  | SPEEX_GET_MODE
  | SPEEX_SET_LOW_MODE
  | SPEEX_GET_LOW_MODE
  | SPEEX_SET_HIGH_MODE
  | SPEEX_GET_HIGH_MODE
  | SPEEX_SET_VBR
  | SPEEX_GET_VBR
  | SPEEX_SET_VBR_QUALITY
  | SPEEX_GET_VBR_QUALITY
  | SPEEX_SET_COMPLEXITY
  | SPEEX_GET_COMPLEXITY
  | SPEEX_SET_BITRATE
  | SPEEX_GET_BITRATE
  | SPEEX_SET_SAMPLING_RATE
  | SPEEX_GET_SAMPLING_RATE
  | SPEEX_SET_VAD
  | SPEEX_GET_VAD
  | SPEEX_SET_ABR
  | SPEEX_GET_ABR
  | SPEEX_SET_DTX
  | SPEEX_GET_DTX
  | SPEEX_SET_SUBMODE_ENCODING
  | SPEEX_GET_SUBMODE_ENCODING
  | SPEEX_SET_PLC_TUNING
  | SPEEX_GET_PLC_TUNING
  | SPEEX_SET_VBR_MAX_BITRATE
  | SPEEX_GET_VBR_MAX_BITRATE
  | SPEEX_SET_HIGHPASS
  | SPEEX_GET_HIGHPASS
  | SPEEX_GET_ACTIVITY

let int_of_control x =
  match x with
    | SPEEX_SET_ENH -> 0
    | SPEEX_GET_ENH -> 1
    | SPEEX_GET_FRAME_SIZE -> 3
    | SPEEX_SET_QUALITY -> 4
    | SPEEX_SET_MODE -> 6
    | SPEEX_GET_MODE -> 7
    | SPEEX_SET_LOW_MODE -> 8
    | SPEEX_GET_LOW_MODE -> 9
    | SPEEX_SET_HIGH_MODE -> 10
    | SPEEX_GET_HIGH_MODE -> 11
    | SPEEX_SET_VBR -> 12
    | SPEEX_GET_VBR -> 13
    | SPEEX_SET_VBR_QUALITY -> 14
    | SPEEX_GET_VBR_QUALITY -> 15
    | SPEEX_SET_COMPLEXITY -> 16
    | SPEEX_GET_COMPLEXITY -> 17
    | SPEEX_SET_BITRATE -> 18
    | SPEEX_GET_BITRATE -> 19
    | SPEEX_SET_SAMPLING_RATE -> 24
    | SPEEX_GET_SAMPLING_RATE -> 25
    | SPEEX_RESET_STATE -> 26
    | SPEEX_SET_VAD -> 30
    | SPEEX_GET_VAD -> 31
    | SPEEX_SET_ABR -> 32
    | SPEEX_GET_ABR -> 33
    | SPEEX_SET_DTX -> 34
    | SPEEX_GET_DTX -> 35
    | SPEEX_SET_SUBMODE_ENCODING -> 36
    | SPEEX_GET_SUBMODE_ENCODING -> 37
    | SPEEX_SET_PLC_TUNING -> 40
    | SPEEX_GET_PLC_TUNING -> 41
    | SPEEX_SET_VBR_MAX_BITRATE -> 42
    | SPEEX_GET_VBR_MAX_BITRATE -> 43
    | SPEEX_SET_HIGHPASS -> 44
    | SPEEX_GET_HIGHPASS -> 45
    | SPEEX_GET_ACTIVITY -> 47

let _ =
  Callback.register "caml_speex_mode_of_int" mode_of_int;
  Callback.register "caml_speex_int_of_mode" int_of_mode;
  Callback.register_exception "ocaml_speex_invfrlen_exn" Invalid_frame_size;
  Callback.register_exception "ocaml_speex_eos_exn" Ogg.End_of_stream

module Header = struct
  type t = {
    id : string;
    version : string;
    version_id : int;
    header_size : int;
    rate : int;
    mode : mode;
    mode_bitstream_version : int;
    nb_channels : int;
    bitrate : int;
    frame_size : int;
    vbr : bool;
    frames_per_packet : int;
    extra_headers : int;
  }

  (* Defined in speex_header.h *)
  let header_string_length = 8
  let header_version_length = 20

  external init : int -> int -> internal_mode -> int -> bool -> t
    = "caml_speex_init_header"

  let init ?(frames_per_packet = 1) ?(mode = Wideband) ?(vbr = true)
      ~nb_channels ~rate () =
    init rate nb_channels (internal_mode_of_mode mode) frames_per_packet vbr

  external encode_header_packetout :
    t -> string array -> Ogg.Stream.packet * Ogg.Stream.packet
    = "caml_speex_encode_header"

  let encode_header_packetout e l =
    let l = List.map (fun (x, y) -> Printf.sprintf "%s=%s" x y) l in
    encode_header_packetout e (Array.of_list l)

  let encode_header e l s =
    let p1, p2 = encode_header_packetout e l in
    Ogg.Stream.put_packet s p1;
    Ogg.Stream.put_packet s p2

  external header_of_packet : Ogg.Stream.packet -> t
    = "caml_speex_header_of_packet"

  external comments_of_packet : Ogg.Stream.packet -> string array
    = "caml_speex_comments_of_packet"

  let comments_of_packet p =
    let x = comments_of_packet p in
    let vendor = x.(0) in
    let x = Array.sub x 1 (Array.length x - 1) in
    let c_k = ref 0 in
    let c_v = ref 0 in
    let split s =
      try
        let i = String.index s '=' in
        try (String.sub s 0 i, String.sub s (i + 1) (String.length s - i - 1))
        with Invalid_argument _ ->
          c_v := !c_v + 1;
          (String.sub s 0 i, Printf.sprintf "unknown_value_%i" !c_k)
      with Not_found ->
        c_k := !c_k + 1;
        (Printf.sprintf "unkown_key_%i" !c_k, s)
    in
    (vendor, Array.to_list (Array.map split x))
end

module Encoder = struct
  type t

  external init : internal_mode -> int -> t = "ocaml_speex_enc_init"

  let init m x = init (internal_mode_of_mode m) x

  external get : t -> int -> int = "ocaml_speex_encoder_ctl_get"

  let get x q = get x (int_of_control q)

  external set : t -> int -> int -> unit = "ocaml_speex_encoder_ctl_set"

  let set x q v = set x (int_of_control q) v

  external encode_page_main :
    t -> int -> Ogg.Stream.stream -> (unit -> float array) -> Ogg.Page.t
    = "ocaml_speex_encode_page"

  let encode_page e s f = encode_page_main e 1 s f

  let merge frame =
    let ret =
      Array.init (2 * Array.length frame.(0)) (fun _ -> frame.(0).(0))
    in
    Array.iteri
      (fun i _ ->
        ret.(2 * i) <- frame.(0).(i);
        ret.((2 * i) + 1) <- frame.(1).(i))
      frame.(0);
    ret

  let encode_page_stereo e s f =
    let f () = merge (f ()) in
    encode_page_main e 2 s f

  external encode_page_int_main :
    t -> int -> Ogg.Stream.stream -> (unit -> int array) -> Ogg.Page.t
    = "ocaml_speex_encode_page_int"

  let encode_page_int e s f = encode_page_int_main e 1 s f

  let encode_page_int_stereo e s f =
    let f () = merge (f ()) in
    encode_page_int_main e 2 s f

  external eos : t -> Ogg.Stream.stream -> unit = "ocaml_speex_encoder_eos"
end

module Decoder = struct
  type t

  external init : internal_mode -> t = "ocaml_speex_dec_init"

  let init m = init (internal_mode_of_mode m)

  external get : t -> int -> int = "ocaml_speex_decoder_ctl_get"

  let get x q = get x (int_of_control q)

  external set : t -> int -> int -> unit = "ocaml_speex_decoder_ctl_set"

  let set x q v = set x (int_of_control q) v

  let split frame =
    let ret = Array.make 2 (Array.make (Array.length frame / 2) frame.(0)) in
    for i = 0 to (Array.length frame / 2) - 1 do
      ret.(0).(i) <- frame.(2 * i);
      ret.(1).(i) <- frame.((2 * i) + 1)
    done;
    ret

  external decode_feed :
    t -> int -> Ogg.Stream.stream -> (float array -> unit) -> unit
    = "ocaml_speex_decoder_decode"

  let decode_gen e s chan func split =
    let l = ref [] in
    let feed x = l := split x :: !l in
    begin
      try func e chan s feed
      with Ogg.Not_enough_data ->
        if List.length !l = 0 then raise Ogg.Not_enough_data
    end;
    List.rev !l

  let decode e s = decode_gen e s 1 decode_feed (fun x -> x)
  let decode_stereo e s = decode_gen e s 2 decode_feed split

  let decode_feed_stereo e s feed =
    let feed x = feed (split x) in
    decode_feed e 2 s feed

  let decode_feed e s feed = decode_feed e 1 s feed

  external decode_int_feed :
    t -> int -> Ogg.Stream.stream -> (int array -> unit) -> unit
    = "ocaml_speex_decoder_decode_int"

  let decode_int e s = decode_gen e s 2 decode_int_feed (fun x -> x)
  let decode_int_stereo e s = decode_gen e s 2 decode_int_feed split

  let decode_int_feed_stereo e s feed =
    let feed x = feed (split x) in
    decode_int_feed e 2 s feed

  let decode_int_feed e s feed = decode_int_feed e 1 s feed
end

module Wrapper = struct
  module Decoder = struct
    exception Not_speex
    exception Internal

    type t =
      (Decoder.t
      * Ogg.Stream.stream
      * Ogg.Sync.t
      * nativeint
      * int
      * (string * (string * string) list)
      * Header.t)
      ref

    type read = bytes -> int -> int -> int

    let open_sync sync =
      (* Test wether the stream contains speex data *)
      let test_speex () =
        (* Get First page *)
        let page = Ogg.Sync.read sync in
        (* Check wether this is a b_o_s *)
        if not (Ogg.Page.bos page) then raise Not_found;
        (* Create a stream with this ID *)
        let serial = Ogg.Page.serialno page in
        let os = Ogg.Stream.create ~serial () in
        Ogg.Stream.put_page os page;
        (* Test header. Do not catch anything, first page should be sufficient *)
        let packet = Ogg.Stream.get_packet os in
        try
          let header = Header.header_of_packet packet in
          (* Get comments *)
          let page = ref (Ogg.Sync.read sync) in
          while Ogg.Page.serialno !page <> serial do
            page := Ogg.Sync.read sync
          done;
          Ogg.Stream.put_page os !page;
          let comments = Ogg.Stream.get_packet os in
          let comments = Header.comments_of_packet comments in
          (serial, os, header, comments)
        with _ ->
          Printf.printf "Not a speex stream..\n";
          raise Internal
      in
      let rec init () =
        try test_speex () with
          (* Not_found is not catched: ogg stream always start
             with all b_o_s and we don't care about sequenced streams here *)
          | Internal -> init ()
          | Ogg.Not_enough_data -> raise Not_speex
      in
      let serial, os, header, comments = init () in
      let chans = header.Header.nb_channels in
      let rate = header.Header.rate in
      let mode = header.Header.mode in
      let dec = Decoder.init mode in
      Decoder.set dec SPEEX_SET_SAMPLING_RATE rate;
      ref (dec, os, sync, serial, chans, comments, header)

    let open_file file =
      let sync, fd = Ogg.Sync.create_from_file file in
      (open_sync sync, fd)

    let open_feed feed =
      let sync = Ogg.Sync.create feed in
      open_sync sync

    let serial x =
      let _, _, _, serial, _, _, _ = !x in
      serial

    let comments x =
      let _, _, _, _, _, (_, l), _ = !x in
      l

    let header x =
      let _, _, _, _, _, _, h = !x in
      h

    let decode_gen f x =
      let dec, os, sync, serial, _, _, _ = !x in
      let dec = ref dec in
      let serial = ref serial in
      let os = ref os in
      let eos = ref false in
      let rec put () =
        try
          if !eos then (
            try
              (* Try to open a new stream *)
              let y = open_sync sync in
              let ndec, nos, _, nserial, _, _, _ = !y in
              x := !y;
              dec := ndec;
              serial := nserial;
              os := nos;
              eos := false
            with Not_found -> raise Internal);
          let page = Ogg.Sync.read sync in
          if Ogg.Stream.eos !os then eos := true;
          if Ogg.Page.serialno page = !serial then Ogg.Stream.put_page !os page
        with
          | Internal | Ogg.Not_enough_data -> raise Ogg.End_of_stream
          | Ogg.End_of_stream ->
              eos := true;
              put ()
      in
      let rec get () =
        try f !dec !os with
          | Ogg.End_of_stream ->
              eos := true;
              put ();
              get ()
          | Ogg.Not_enough_data ->
              put ();
              get ()
      in
      get ()

    let decode v = decode_gen Decoder.decode v
    let decode_stereo v = decode_gen Decoder.decode_stereo v

    let decode_feed v feed =
      decode_gen (fun x y -> Decoder.decode_feed x y feed) v

    let decode_feed_stereo v feed =
      decode_gen (fun x y -> Decoder.decode_feed_stereo x y feed) v

    let decode_int v = decode_gen Decoder.decode_int v
    let decode_int_stereo v = decode_gen Decoder.decode_int_stereo v

    let decode_int_feed v feed =
      decode_gen (fun x y -> Decoder.decode_int_feed x y feed) v

    let decode_int_feed_stereo v feed =
      decode_gen (fun x y -> Decoder.decode_int_feed_stereo x y feed) v
  end
end

module Skeleton = struct
  external fisbone :
    nativeint -> Header.t -> Int64.t -> string -> Ogg.Stream.packet
    = "ocaml_speex_skeleton_fisbone"

  let fisbone ?(start_granule = Int64.zero)
      ?(headers = [("Content-type", "audio/speex")]) ~serialno ~header () =
    let concat s (h, v) = Printf.sprintf "%s%s: %s\r\n" s h v in
    let s = List.fold_left concat "" headers in
    fisbone serialno header start_granule s
end