()
| 32 | |
| 33 | |
| 34 | def main(): |
| 35 | parser = argparse.ArgumentParser(description="Generate subtitles from audio/video using FunASR") |
| 36 | parser.add_argument("input", help="Audio/video file path") |
| 37 | parser.add_argument("-o", "--output", help="Output file (default: input.srt)") |
| 38 | parser.add_argument("--format", choices=["srt", "vtt"], default="srt") |
| 39 | parser.add_argument("--model", default="iic/SenseVoiceSmall") |
| 40 | parser.add_argument("--device", default="cuda") |
| 41 | parser.add_argument("--spk", action="store_true", help="Include speaker labels") |
| 42 | parser.add_argument("--lang", default="auto") |
| 43 | args = parser.parse_args() |
| 44 | |
| 45 | if not os.path.exists(args.input): |
| 46 | print(f"Error: {args.input} not found") |
| 47 | sys.exit(1) |
| 48 | |
| 49 | output_path = args.output or f"{os.path.splitext(args.input)[0]}.{args.format}" |
| 50 | print(f"Input: {args.input}") |
| 51 | print(f"Output: {output_path}") |
| 52 | |
| 53 | from funasr import AutoModel |
| 54 | |
| 55 | kwargs = {"model": args.model, "vad_model": "fsmn-vad", |
| 56 | "vad_kwargs": {"max_single_segment_time": 30000}, |
| 57 | "device": args.device, "disable_update": True} |
| 58 | if args.spk: |
| 59 | kwargs["spk_model"] = "cam++" |
| 60 | if "Fun-ASR-Nano" in args.model or "Qwen" in args.model: |
| 61 | kwargs["trust_remote_code"] = True |
| 62 | kwargs["hub"] = "hf" |
| 63 | |
| 64 | print("Loading model...") |
| 65 | model = AutoModel(**kwargs) |
| 66 | print("Transcribing...") |
| 67 | result = model.generate(input=args.input, batch_size=1) |
| 68 | |
| 69 | segments = [] |
| 70 | if "sentence_info" in result[0]: |
| 71 | for seg in result[0]["sentence_info"]: |
| 72 | text = re.sub(r'<\|[^|]*\|>', '', seg.get("text", "")).strip() |
| 73 | if text: |
| 74 | segments.append({"start": seg.get("start", 0), "end": seg.get("end", 0), |
| 75 | "text": text, "spk": seg.get("spk")}) |
| 76 | else: |
| 77 | text = re.sub(r'<\|[^|]*\|>', '', result[0].get("text", "")).strip() |
| 78 | if text: |
| 79 | segments.append({"start": 0, "end": 0, "text": text, "spk": None}) |
| 80 | |
| 81 | if not segments: |
| 82 | print("No speech detected.") |
| 83 | sys.exit(0) |
| 84 | |
| 85 | fmt = format_time_srt if args.format == "srt" else format_time_vtt |
| 86 | with open(output_path, "w", encoding="utf-8") as f: |
| 87 | if args.format == "vtt": |
| 88 | f.write("WEBVTT\n\n") |
| 89 | for i, seg in enumerate(segments, 1): |
| 90 | text = f"[Speaker {seg['spk']}] {seg['text']}" if args.spk and seg['spk'] is not None else seg['text'] |
| 91 | if args.format == "srt": |
no test coverage detected
searching dependent graphs…