From 7c915710b8778b1b73235299187a7aa46ad0ebcc Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Fri, 6 Jun 2025 16:53:30 -0400 Subject: [PATCH 01/21] type checking, beginning of unit tests --- .github/workflows/ci.yml | 58 ++ output_captions.srt | 812 ------------------ pyproject.toml | 37 + requirements.txt | 2 +- setup.cfg | 4 + src/__init__.py | 0 src/captionalchemy/__init__.py | 1 + src/captionalchemy/__main__.py | 4 + .../caption.py} | 137 ++- .../tools/audio_analysis/__init__.py | 0 .../audio_segment_integration.py | 3 +- .../tools/audio_analysis/diarization.py | 0 .../audio_analysis/non_speech_detection.py | 5 +- .../tools/audio_analysis/vad.py | 0 .../tools/captioning/timing_analyzer.py | 16 +- .../tools/captioning/transcriber.py | 4 +- .../tools/captioning/words_dictionary.json | 0 .../tools/captioning/writers/sami_writer.py | 212 +++++ .../tools/captioning/writers/srt_writer.py} | 0 .../tools/captioning/writers/vtt_writer.py | 118 +++ .../tools/cv}/embed_known_faces.py | 13 +- .../tools/cv}/recognize_faces.py | 0 .../tools/media_utils}/download_video.py | 1 - .../tools/media_utils}/extract_audio.py | 0 src/tools/audio_analysis/__init__.py | 0 .../captioning/writers/vtt_caption_writer.py | 0 tests/conftest.py | 251 ++++++ 27 files changed, 824 insertions(+), 854 deletions(-) create mode 100644 .github/workflows/ci.yml delete mode 100644 output_captions.srt create mode 100644 setup.cfg delete mode 100644 src/__init__.py create mode 100644 src/captionalchemy/__init__.py create mode 100644 src/captionalchemy/__main__.py rename src/{inference.py => captionalchemy/caption.py} (58%) rename __init__.py => src/captionalchemy/tools/audio_analysis/__init__.py (100%) rename src/{ => captionalchemy}/tools/audio_analysis/audio_segment_integration.py (99%) rename src/{ => captionalchemy}/tools/audio_analysis/diarization.py (100%) rename src/{ => captionalchemy}/tools/audio_analysis/non_speech_detection.py (94%) rename src/{ => captionalchemy}/tools/audio_analysis/vad.py (100%) rename src/{ => captionalchemy}/tools/captioning/timing_analyzer.py (97%) rename src/{ => captionalchemy}/tools/captioning/transcriber.py (98%) rename src/{ => captionalchemy}/tools/captioning/words_dictionary.json (100%) create mode 100644 src/captionalchemy/tools/captioning/writers/sami_writer.py rename src/{tools/captioning/writers/srt_caption_writer.py => captionalchemy/tools/captioning/writers/srt_writer.py} (100%) create mode 100644 src/captionalchemy/tools/captioning/writers/vtt_writer.py rename src/{tools => captionalchemy/tools/cv}/embed_known_faces.py (80%) rename src/{tools => captionalchemy/tools/cv}/recognize_faces.py (100%) rename src/{tools => captionalchemy/tools/media_utils}/download_video.py (99%) rename src/{tools => captionalchemy/tools/media_utils}/extract_audio.py (100%) delete mode 100644 src/tools/audio_analysis/__init__.py delete mode 100644 src/tools/captioning/writers/vtt_caption_writer.py create mode 100644 tests/conftest.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..65fca49 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,58 @@ +name: CI + +on: + push: + tags: + - "v*.*.*" + +jobs: + build_and_test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: [3.8, 3.9, 3.10, 3.11] + steps: + - name: Checkout code + uses: actions/checkout@v3 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + - name: Lint with flake8 + run: | + python -m pip install --upgrade flake8 + flake8 src/ tests/ + - name: Run type checks with mypy + run: | + python -m pip install --upgrade mypy + mypy --install-types + mypy src/ + - name: Run tests + run: | + + publish: + needs: build_and_test + runs-on: ubuntu-latest + if: startsWith(github.ref, 'refs/tags/v') + steps: + - name: Checkout code + uses: actions/checkout@v3 + - name: Set up Python 3.11.5 + uses: actions/setup-python@v4 + with: + python-version: 3.11.5 + - name: Build distributions + run: | + python -m pip install --upgrade build + python -m build + - name: Publish to PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: | + python -m pip install --upgrade twine + twine upload --skip-existing dist/* diff --git a/output_captions.srt b/output_captions.srt deleted file mode 100644 index a32aa9f..0000000 --- a/output_captions.srt +++ /dev/null @@ -1,812 +0,0 @@ -1 -00:00:00,000 --> 00:00:02,560 -[MUSIC PLAYING] - -2 -00:00:02,560 --> 00:00:03,254 -[SILENCE] - -3 -00:00:03,294 --> 00:00:07,314 -SpaceTime Dude: If you have perfect knowledge of every single particle in the universe, can - -4 -00:00:07,414 --> 00:00:11,454 -you use the laws of physics to rewind all the way back to the Big - -5 -00:00:11,554 --> 00:00:15,474 -Bang? Is the entire history of the universe perfectly knowable - -6 -00:00:15,359 --> 00:00:16,640 -[MUSIC PLAYING] - -7 -00:00:15,574 --> 00:00:16,524 -? - -8 -00:00:16,192 --> 00:00:19,472 -or has information somehow been lost along the way? - -9 -00:00:19,199 --> 00:00:25,280 -[MUSIC PLAYING] - -10 -00:00:25,491 --> 00:00:28,481 -The laws of physics are equations of motion. They - -11 -00:00:28,581 --> 00:00:33,181 -are mathematical rules that dictate how systems evolve in time. Newton's - -12 -00:00:33,281 --> 00:00:35,471 -equations for classical mechanics, maxwell's - -13 -00:00:35,571 --> 00:00:37,291 -equations for electromagnetism, and - -14 -00:00:37,391 --> 00:00:40,591 -the Schrodinger equation for quantum mechanics. These - -15 -00:00:40,691 --> 00:00:45,231 -laws can be used to predict how the universe will evolve into the future. They - -16 -00:00:45,331 --> 00:00:47,231 -are deterministic. Perfect - -17 -00:00:47,331 --> 00:00:49,381 -knowledge of a system in the present, perfectly - -18 -00:00:49,481 --> 00:00:52,691 -predicts how the system will change in the next instant, and - -19 -00:00:52,791 --> 00:00:54,141 -the instant after that, add - -20 -00:00:54,241 --> 00:00:59,851 -infinitum. But determinism in the forward time direction does not guarantee - -21 -00:00:59,951 --> 00:01:03,321 -that the same laws can perfectly predict the past. And - -22 -00:01:03,421 --> 00:01:05,931 -yet this sort of deterministic symmetry, time - -23 -00:01:06,031 --> 00:01:07,751 -reversal symmetry, is - -24 -00:01:07,851 --> 00:01:11,401 -essential for information itself to be conserved. - -25 -00:01:11,547 --> 00:01:12,053 -[SILENCE] - -26 -00:01:12,133 --> 00:01:16,203 -Today, we learn why conservation of information is such a - -27 -00:01:16,303 --> 00:01:19,293 -fundamental requirement of quantum mechanics. In - -28 -00:01:19,393 --> 00:01:20,923 -a future episode, we'll - -29 -00:01:21,023 --> 00:01:24,483 -see how this law might be broken by black holes. An - -30 -00:01:24,583 --> 00:01:28,603 -important tool for making predictions with the laws of physics are the conservation laws - -31 -00:01:28,703 --> 00:01:30,463 -that we derive from them. As - -32 -00:01:30,563 --> 00:01:31,883 -we discussed recently, we - -33 -00:01:31,983 --> 00:01:36,023 -can use Notice Theorem to discover quantities that are conserved in a system by - -34 -00:01:36,123 --> 00:01:39,193 -looking at the symmetries of the equations of motion. - -35 -00:01:39,790 --> 00:01:43,970 -For example, if the equations stay the same from one point in time to the next, then - -36 -00:01:44,070 --> 00:01:45,260 -energy is conserved. - -37 -00:01:45,612 --> 00:01:48,732 -Notice theorem applies to smooth, continuous - -38 -00:01:48,832 --> 00:01:52,142 -symmetries where we can transform the chosen coordinate, however - -39 -00:01:52,242 --> 00:01:54,152 -much we like without changing the system. - -40 -00:01:54,857 --> 00:01:59,337 -but there are also symmetries under discrete transformations. They're - -41 -00:01:59,437 --> 00:02:01,307 -more like on/off switches. For - -42 -00:02:01,407 --> 00:02:05,407 -example, we could reverse all electric charges or we could flip the x -axis - -43 -00:02:05,507 --> 00:02:08,547 -by looking in a mirror or we could make time run backwards. - -44 -00:02:09,292 --> 00:02:10,972 -Symmetry in that last one, time - -45 -00:02:11,072 --> 00:02:12,632 -reversal symmetry, isn't - -46 -00:02:12,732 --> 00:02:14,362 -covered by notice theorem, but - -47 -00:02:14,462 --> 00:02:16,732 -it's still tied to a conservation law, the - -48 -00:02:16,832 --> 00:02:19,342 -conservation of information. We - -49 -00:02:19,442 --> 00:02:23,552 -say that a system is time reversal symmetric if its equations of motion allow - -50 -00:02:23,652 --> 00:02:27,652 -us to perfectly predict the starting point simply by knowing the state of - -51 -00:02:27,752 --> 00:02:29,372 -the system at any later time. If - -52 -00:02:29,472 --> 00:02:33,542 -they allow us to rewind the clock and figure out a single unique - -53 -00:02:33,642 --> 00:02:34,592 -history, - -54 -00:02:34,369 --> 00:02:38,709 -The entire universe would be time reversal symmetric if knowing the exact - -55 -00:02:38,809 --> 00:02:43,309 -state of every particle in the universe at one point in time allowed us to calculate - -56 -00:02:43,409 --> 00:02:46,239 -its exact past history at all times. That - -57 -00:02:46,339 --> 00:02:50,719 -would mean that the exact configuration of the universe at any point in time defines - -58 -00:02:50,819 --> 00:02:52,769 -the exact configuration at any other point. - -59 -00:02:53,537 --> 00:02:54,547 -For one thing, this - -60 -00:02:54,647 --> 00:02:58,727 -time reversal symmetry means complete information about the configuration at all - -61 -00:02:58,827 --> 00:03:00,847 -prior times still exists, and - -62 -00:03:00,947 --> 00:03:04,427 -always will, even if we can't practically access it. - -63 -00:03:04,969 --> 00:03:07,879 -That's what I mean about conservation of information. - -64 -00:03:08,426 --> 00:03:11,436 -A related idea is causal determinism, the - -65 -00:03:11,536 --> 00:03:15,616 -idea that perfectly knowing the current state perfectly predicts all future - -66 -00:03:15,716 --> 00:03:19,796 -states. But this sort of future determinism doesn't have to be time - -67 -00:03:19,896 --> 00:03:21,376 -reversal symmetric. It's - -68 -00:03:21,476 --> 00:03:25,386 -possible for the future to be perfectly predictable by the laws of physics while the - -69 -00:03:25,486 --> 00:03:26,436 -past is not. - -70 -00:03:26,750 --> 00:03:30,770 -For example, what if many different configurations of particles in the present could - -71 -00:03:30,870 --> 00:03:34,100 -converge on a single configuration of particles in the future? - -72 -00:03:34,575 --> 00:03:37,925 -If a multitude of states can all evolve into the same state, then - -73 -00:03:38,025 --> 00:03:41,825 -knowing the latter state isn't enough for us to figure out what past states led to it. - -74 -00:03:42,505 --> 00:03:43,535 -For example, imagine - -75 -00:03:43,635 --> 00:03:44,955 -you have two states, a - -76 -00:03:45,055 --> 00:03:49,205 -and B, and the laws of motion say that A becomes C and B - -77 -00:03:49,305 --> 00:03:53,265 -becomes D. If we look at the later state and C states C - -78 -00:03:53,365 --> 00:03:55,205 -, we know A came first, but - -79 -00:03:55,305 --> 00:03:56,575 -if we see state D, we - -80 -00:03:56,675 --> 00:03:58,695 -know B came first. But - -81 -00:03:58,795 --> 00:04:01,375 -what if both A and B lead to state C? - -82 -00:04:01,275 --> 00:04:01,832 -[SILENCE] - -83 -00:04:01,952 --> 00:04:04,752 -that future direction is deterministic. But - -84 -00:04:04,852 --> 00:04:06,502 -now, if we look at the final state, we - -85 -00:04:06,602 --> 00:04:10,892 -see state c, but we don't know which of a or b led to it. - -86 -00:04:10,992 --> 00:04:11,407 -The - -87 -00:04:11,507 --> 00:04:13,752 -information about those past states would be lost. - -88 -00:04:14,322 --> 00:04:14,942 -Okay, so - -89 -00:04:15,042 --> 00:04:17,342 -it should be simple enough to erase information, right - -90 -00:04:17,442 --> 00:04:21,432 -? We just set things up so our laws of motion force two possible initial - -91 -00:04:21,532 --> 00:04:24,772 -states into the same exact final state. Then - -92 -00:04:24,872 --> 00:04:27,152 -we wouldn't know what the original state was, and - -93 -00:04:27,252 --> 00:04:29,002 -information would be destroyed. - -94 -00:04:29,583 --> 00:04:30,773 -Actually, quantum - -95 -00:04:30,873 --> 00:04:32,803 -mechanics forbids this. It - -96 -00:04:32,903 --> 00:04:38,023 -ensures conservation of information and time reversal symmetry because of an even more fundamental - -97 -00:04:38,123 --> 00:04:40,883 -rule. The conservation of probability. - -98 -00:04:40,898 --> 00:04:41,505 -[SILENCE] - -99 -00:04:41,575 --> 00:04:42,515 -To get at this, let's - -100 -00:04:42,615 --> 00:04:45,825 -think about the basic equation of motion of quantum mechanics. That's - -101 -00:04:45,925 --> 00:04:47,795 -the Schrodinger equation. The - -102 -00:04:47,895 --> 00:04:51,855 -time -dependent Schrodinger equation describes the time evolution of this - -103 -00:04:51,955 --> 00:04:53,695 -thing called the wavefunction. The - -104 -00:04:53,795 --> 00:04:57,735 -wavefunction of a system fully describes all of its properties, and - -105 -00:04:57,835 --> 00:04:59,685 -in quantum mechanics, that - -106 -00:04:59,785 --> 00:05:02,885 -means the probability distribution of all of its properties, which - -107 -00:05:02,985 --> 00:05:05,505 -you can get by taking the square of the wavefunction. - -108 -00:05:05,907 --> 00:05:06,447 -[SILENCE] - -109 -00:05:06,567 --> 00:05:10,557 -For example, the wave function of a particle encapsulates the probability - -110 -00:05:10,657 --> 00:05:14,597 -that will be found in this or that location if we were to try to measure - -111 -00:05:14,697 --> 00:05:18,607 -its position. The Schrödinger equation perfectly predicts both - -112 -00:05:18,707 --> 00:05:23,397 -the past and future evolution of a given wave function in any given environment - -113 -00:05:23,497 --> 00:05:24,937 -, or in quantum speak, in - -114 -00:05:25,037 --> 00:05:26,447 -any given potential. - -115 -00:05:26,465 --> 00:05:27,465 -In that sense, it's - -116 -00:05:27,565 --> 00:05:30,495 -deterministic and time reversal symmetric, and - -117 -00:05:30,595 --> 00:05:32,065 -so conserves information. - -118 -00:05:32,485 --> 00:05:36,595 -in principle, a given wave function in a given potential could mean - -119 -00:05:36,695 --> 00:05:40,125 -the wave function of an electron moving in an atom's electric field, or - -120 -00:05:40,225 --> 00:05:44,535 -it could mean the wave function of the entire universe in its own impossibly - -121 -00:05:44,635 --> 00:05:46,645 -complex and changing potential. - -122 -00:05:46,845 --> 00:05:47,487 -[SILENCE] - -123 -00:05:47,567 --> 00:05:52,017 -So, the Schrodinger equation guarantees time reversibility and so the conservation - -124 -00:05:52,117 --> 00:05:53,237 -of information, as - -125 -00:05:53,337 --> 00:05:57,337 -do the more advanced formulations of quantum mechanics like the Dirac equation and quantum - -126 -00:05:57,437 --> 00:05:58,387 -field theories. - -127 -00:05:58,506 --> 00:06:03,726 -that guarantee arises from a really fundamental foundational quality of these theories - -128 -00:06:03,826 --> 00:06:06,686 -. It comes from what we call unitarity. Remember - -129 -00:06:06,786 --> 00:06:11,296 -that the wave function encapsulates the distribution of probabilities for a given property. By - -130 -00:06:11,396 --> 00:06:12,686 -definition, those - -131 -00:06:12,786 --> 00:06:15,526 -probabilities should add up to one. That - -132 -00:06:15,626 --> 00:06:20,686 -just means there's a 100% probability that any given property will have some possible value - -133 -00:06:20,786 --> 00:06:22,426 -, even if that value is zero. - -134 -00:06:23,102 --> 00:06:25,462 -In the case of particle position, probability - -135 -00:06:25,562 --> 00:06:29,022 -of adding to one just means that the particle is definitely somewhere. - -136 -00:06:29,761 --> 00:06:31,171 -And as time goes on, a - -137 -00:06:31,271 --> 00:06:35,191 -particle's properties will continue to have possible values. The - -138 -00:06:35,291 --> 00:06:37,691 -probability should continue to sum to 1. - -139 -00:06:37,875 --> 00:06:38,500 -[SILENCE] - -140 -00:06:38,590 --> 00:06:39,980 -If this is true, and - -141 -00:06:40,080 --> 00:06:44,920 -it must be, we say that the time evolution of the wave function is unitary - -142 -00:06:45,020 --> 00:06:49,010 -. And this unitarity is a foundational assumption in all formulations - -143 -00:06:49,110 --> 00:06:51,520 -of quantum mechanics and quantum field theories. - -144 -00:06:52,066 --> 00:06:56,656 -Uniterity is a non -negotiable statement about how probability works, but - -145 -00:06:56,756 --> 00:07:01,086 -the condition also ensures time reversal symmetry and conservation of information. - -146 -00:07:01,600 --> 00:07:05,340 -It's not straightforward to explain without getting into some hairy math, but - -147 -00:07:05,440 --> 00:07:09,790 -the upshot is that quantum states must remain independent of each other - -148 -00:07:09,890 --> 00:07:12,380 -in order to preserve probability. Two - -149 -00:07:12,480 --> 00:07:16,610 -independent quantum states can't evolve into the exact same quantum - -150 -00:07:16,710 --> 00:07:20,190 -state. If they did then the probabilities for the initial states, or - -151 -00:07:20,290 --> 00:07:22,940 -for the final state can't both sum to one. - -152 -00:07:23,855 --> 00:07:25,725 -Back to our first example, quantum - -153 -00:07:25,825 --> 00:07:29,225 -states A and B can't both become quantum states C. The - -154 -00:07:29,325 --> 00:07:33,645 -sum of the probabilities prior to the merger would not equal the sum of the probabilities - -155 -00:07:33,745 --> 00:07:37,595 -after the merger and unitarity would be broken. The - -156 -00:07:37,695 --> 00:07:42,085 -only type of evolution that preserves probability and unitarity is the evolution - -157 -00:07:42,185 --> 00:07:45,365 -that also preserves the number of quantum states. And - -158 -00:07:45,465 --> 00:07:49,495 -the preservation of quantum states means preservation of information because you - -159 -00:07:49,595 --> 00:07:53,545 -can trace a quantum state indefinitely forwards and backwards in - -160 -00:07:53,645 --> 00:07:57,625 -time. But all of this talk of quantum mechanics being deterministic seems - -161 -00:07:57,725 --> 00:08:01,725 -a bit at odds with the idea of quantum randomness and the uncertainty principle - -162 -00:08:01,825 --> 00:08:02,835 -. After all, doesn't - -163 -00:08:02,935 --> 00:08:06,945 -the act of measurement pick a single value of some quantity from the range - -164 -00:08:07,045 --> 00:08:08,535 -of possible values? That - -165 -00:08:08,635 --> 00:08:12,605 -value seems to be chosen randomly based on the probability distribution - -166 -00:08:12,705 --> 00:08:14,565 -encoded in the wave function. And - -167 -00:08:14,665 --> 00:08:18,755 -the precision of the knowability of that value is defined by the uncertainty - -168 -00:08:18,855 --> 00:08:21,895 -principle. It sure seems like information can be lost. - -169 -00:08:22,385 --> 00:08:25,915 -but that's not really what we mean by quantum information. Quantum - -170 -00:08:26,015 --> 00:08:29,945 -information refers to the full information content of the wave function - -171 -00:08:30,045 --> 00:08:31,845 -, not just what we measure. And - -172 -00:08:31,945 --> 00:08:33,005 -in principle, make - -173 -00:08:33,105 --> 00:08:37,225 -enough measurements and you can extract all of the information from a wave function. - -174 -00:08:37,325 --> 00:08:37,825 -It's - -175 -00:08:37,925 --> 00:08:41,705 -worth mentioning that the collapse of the wave function in the Copenhagen interpretation - -176 -00:08:41,805 --> 00:08:45,985 -of quantum mechanics actually does mess with conservation of information - -177 -00:08:46,085 --> 00:08:48,125 -. According to that interpretation, the - -178 -00:08:48,225 --> 00:08:52,345 -active measurement actually alters the entire wave function, causing - -179 -00:08:52,445 --> 00:08:56,655 -it to shrink down to the narrow range of possible values implied by the - -180 -00:08:56,755 --> 00:08:58,475 -result of the measurement. - -181 -00:08:58,754 --> 00:09:04,244 -but that measured wave function can't then be tracked backwards to recover the passed wave function - -182 -00:09:04,344 --> 00:09:09,484 -. So this interpretation is neither deterministic nor time reversal symmetric - -183 -00:09:09,584 --> 00:09:10,534 -. - -184 -00:09:10,005 --> 00:09:12,165 -Other interpretations of quantum mechanics, for - -185 -00:09:12,265 --> 00:09:13,090 -example, everett's - -186 -00:09:13,190 --> 00:09:16,455 -many worlds or the de Broglie -Bohm pilot wave theory, preserve - -187 -00:09:16,555 --> 00:09:18,485 -this time reversibility. In - -188 -00:09:18,585 --> 00:09:20,065 -the case of many worlds, the - -189 -00:09:20,165 --> 00:09:23,445 -entire wave function continues to exist even after measurement. - -190 -00:09:23,900 --> 00:09:28,270 -That measurement itself is just a thing that happens in one part of the possibility space - -191 -00:09:28,370 --> 00:09:31,160 -, no information is really lost. And - -192 -00:09:31,260 --> 00:09:32,950 -in the case of pilot wave theory, the - -193 -00:09:33,050 --> 00:09:37,340 -wave function contains hidden information that is carried with the final measured - -194 -00:09:37,440 --> 00:09:38,390 -particle. - -195 -00:09:37,881 --> 00:09:38,877 -[SILENCE] - -196 -00:09:38,937 --> 00:09:43,427 -There is one situation where time reversibility appears to be broken, regardless - -197 -00:09:43,527 --> 00:09:46,457 -of your favourite interpretation of quantum mechanics. That's - -198 -00:09:46,557 --> 00:09:50,417 -the case of black holes and Hawking radiation. Stephen - -199 -00:09:50,517 --> 00:09:55,267 -Hawking's eponymous radiation appears to destroy quantum information, leading - -200 -00:09:55,367 --> 00:09:58,587 -to the famous black hole information paradox. In - -201 -00:09:58,687 --> 00:10:00,087 -an upcoming episode, we'll - -202 -00:10:00,187 --> 00:10:05,427 -see where the quantum information really can be deleted from the otherwise perfect memory - -203 -00:10:05,527 --> 00:10:06,477 -of spacetime. - diff --git a/pyproject.toml b/pyproject.toml index e69de29..4d9f532 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -0,0 +1,37 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "captionalchemy" +version = "0.1.0" +description = "A Python package to create closed captions with face detection and recognition." +readme = "README.md" +license = { text = "MIT" } +authors = [{ name = "Ben Batman", email = "benbatman2@gmail.com" }] + +dependencies = [ + "pyannote.audio==3.3.2", + "pydub==0.25.1", + "openai-whisper==20240930", + "python-dotenv==1.1.0", + "opencv-python==4.11.0.86", + "numpy==2.1.3", + "protobuf==5.29.4", + "insightface==0.7.3", + "onnxruntime==1.22.0", + "panns-inference==0.1.1", + "librosa==0.11.0", + "tqdm==4.67.1", +] + +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", +] + +requires-python = ">=3.8" + +[project.scripts] +captionalchemy = "captionalchemy.caption:main" diff --git a/requirements.txt b/requirements.txt index e86fdae..06507c0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,10 +3,10 @@ pydub==0.25.1 openai-whisper==20240930 python-dotenv==1.1.0 opencv-python==4.11.0.86 -retina-face==0.0.17 numpy==2.1.3 protobuf==5.29.4 insightface==0.7.3 onnxruntime==1.22.0 panns-inference==0.1.1 librosa==0.11.0 +tqdm==4.67.1 diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..3d0d248 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,4 @@ +[flake8] +max-line-length = 110 +extend-ignore = E203, W503 +exclude = [".git", "__pycache__", "build", "dist", ".venv"] diff --git a/src/__init__.py b/src/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/captionalchemy/__init__.py b/src/captionalchemy/__init__.py new file mode 100644 index 0000000..3dc1f76 --- /dev/null +++ b/src/captionalchemy/__init__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/src/captionalchemy/__main__.py b/src/captionalchemy/__main__.py new file mode 100644 index 0000000..8b2897c --- /dev/null +++ b/src/captionalchemy/__main__.py @@ -0,0 +1,4 @@ +from captionalchemy.caption import main + +if __name__ == "__main__": + main() diff --git a/src/inference.py b/src/captionalchemy/caption.py similarity index 58% rename from src/inference.py rename to src/captionalchemy/caption.py index a7f61f7..6eae8fe 100644 --- a/src/inference.py +++ b/src/captionalchemy/caption.py @@ -2,6 +2,8 @@ import tempfile import os import uuid +from typing import Literal +import argparse from dotenv import load_dotenv, find_dotenv from tqdm import tqdm @@ -9,31 +11,50 @@ import torch from tools.audio_analysis.diarization import diarize -from tools.download_video import VideoManager -from tools.recognize_faces import recognize_faces -from tools.extract_audio import extract_audio +from captionalchemy.tools.media_utils.download_video import VideoManager +from captionalchemy.tools.cv.recognize_faces import recognize_faces +from captionalchemy.tools.media_utils.extract_audio import extract_audio from tools.captioning.transcriber import Transcriber from tools.captioning.timing_analyzer import TimingAnalyzer -from tools.embed_known_faces import embed_faces -from tools.captioning.writers.srt_caption_writer import SRTCaptionWriter +from captionalchemy.tools.cv.embed_known_faces import embed_faces +from captionalchemy.tools.captioning.writers.srt_writer import SRTCaptionWriter +from captionalchemy.tools.captioning.writers.vtt_writer import VTTCaptionWriter +from captionalchemy.tools.captioning.writers.sami_writer import SAMICaptionWriter from tools.audio_analysis.vad import get_speech_segments from tools.audio_analysis.non_speech_detection import detect_non_speech_segments from tools.audio_analysis.audio_segment_integration import ( integrate_audio_segments, ) +logger = logging.getLogger(__name__) -def main( + +def run_pipeline( video_url_or_path: str, character_identification: bool = True, known_faces_json: str = "example/known_faces.json", embed_faces_json: str = "example/embed_faces.json", + caption_output_path: str = "output_captions", + caption_format: Literal["vtt", "srt", "smi"] = "srt", ): - """Main function to run the inference pipeline.""" + """ + Core pipeline that: + 1. Embeds known faces (if enabled), + 2. Downloads/extracts audio, + 3. Runs VAD + diarization, + 4. Runs Whisper transcription + face-based speaker ID, + 5. Writes captions in the chosen format (SRT, VTT, SMI). + """ logger.info("Embedding known faces...") embed_faces(known_faces_json, embed_faces_json) video_manager = VideoManager(use_file_buffer=False) - writer = SRTCaptionWriter() + if caption_format == "srt": + writer = SRTCaptionWriter() + elif caption_format == "vtt": + writer = VTTCaptionWriter() + elif caption_format == "smi": + writer = SAMICaptionWriter() + transcriber = Transcriber() timing_analyzer = TimingAnalyzer() @@ -72,10 +93,12 @@ def main( return # Diarize - # diarization_result = diarize(audio_path) looks like this: { "SPEAKER_00": {"start": 3.25409375, "end": 606.2990937500001}, ..., SPEAKER_XX: {} } - diarization_result = { - "SPEAKER_00": {"start": 3.25409375, "end": 606.2990937500001} - } + diarization_result = diarize( + audio_path + ) # { "SPEAKER_00": {"start": 3.25409375, "end": 606.2990937500001}, ..., } + # diarization_result = { + # "SPEAKER_00": {"start": 3.25409375, "end": 606.2990937500001} + # } logger.info("Completed diarization.") logger.debug(f"Diarization result: {diarization_result}") @@ -87,9 +110,7 @@ def main( diarization_result, total_audio_duration=None, ) - - print("Integrated audio events:\n") - print(integrated_audio_events) + logger.debug(f"Integrated audio events: {integrated_audio_events}") # Run whisper and character identification on each individual speaker segment for audio_event in tqdm( @@ -161,18 +182,88 @@ def main( ) # Write the captions to an SRT file - writer.write("output_captions.srt") + writer.write(f"{caption_output_path}.{caption_format}") logger.info("Captions written to output_captions.srt") -if __name__ == "__main__": - logging.basicConfig(level=logging.INFO) - logger = logging.getLogger(__name__) +def _build_arg_parser() -> argparse.ArgumentParser: + """Constructs the CLI argument parser for captionalchemy.""" + parser = argparse.ArgumentParser( + prog="captionalchemy", + description="Download/extract audio from a video," + "run diarization + transcription, and output captions.", + ) + + parser.add_argument( + "video", + help="URL or local path of the video to caption (e.g., 'myvideo.mp4' or 'https://...').", + ) + + parser.add_argument( + "-f", + "--format", + choices=["srt", "vtt", "smi"], + default="srt", + help="Caption output format (default: 'srt').", + ) + + parser.add_argument( + "-o", + "--output", + default="output_captions", + help="Base path (without extension) for output captions. " + "Actual file will be written as `.` (default: 'output_captions').", + ) + + parser.add_argument( + "--no-face-id", + action="store_false", + dest="character_identification", + help="Disable character identification via face recognition.", + ) + parser.add_argument( + "--known-faces-json", + default="example/known_faces.json", + help="Path to JSON file listing known faces to embed (default: 'example/known_faces.json').", + ) + + parser.add_argument( + "--embed-faces-json", + default="example/embed_faces.json", + help="JSON path to store face embeddings (default: 'example/embed_faces.json').", + ) + + parser.add_argument( + "-v", "--verbose", action="store_true", help="Enable debug logging output." + ) + + return parser + + +def main(): load_dotenv(find_dotenv(), override=True) + parser = _build_arg_parser() + args = parser.parse_args() + + if args.verbose: + logging.getLogger().setLevel(logging.DEBUG) + else: + logging.getLogger().setLevel(logging.INFO) - video_url = "mfmt0hf5_spac_426_fordotorg-mp4-720p-3000k.mp4" - logger.info("Starting inference pipeline...") - main( - video_url, + logger = logging.getLogger("captionalchemy") + logger.info("Starting captionalchemy pipeline...") + run_pipeline( + video_url_or_path=args.video, + character_identification=args.character_identification, + known_faces_json=args.known_faces_json, + embed_faces_json=args.embed_faces_json, + caption_output_path=args.output, + caption_format=args.format, ) + logger.info("Pipeline completed successfully.") + logger.info("Output captions can be found at: %s.%s", args.output, args.format) + + +if __name__ == "__main__": + main() diff --git a/__init__.py b/src/captionalchemy/tools/audio_analysis/__init__.py similarity index 100% rename from __init__.py rename to src/captionalchemy/tools/audio_analysis/__init__.py diff --git a/src/tools/audio_analysis/audio_segment_integration.py b/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py similarity index 99% rename from src/tools/audio_analysis/audio_segment_integration.py rename to src/captionalchemy/tools/audio_analysis/audio_segment_integration.py index 05b54d1..8450f8c 100644 --- a/src/tools/audio_analysis/audio_segment_integration.py +++ b/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py @@ -1,5 +1,4 @@ -from typing import List, Dict, Any, Optional, Tuple -import logging +from typing import List, Dict, Optional from dataclasses import dataclass from enum import Enum diff --git a/src/tools/audio_analysis/diarization.py b/src/captionalchemy/tools/audio_analysis/diarization.py similarity index 100% rename from src/tools/audio_analysis/diarization.py rename to src/captionalchemy/tools/audio_analysis/diarization.py diff --git a/src/tools/audio_analysis/non_speech_detection.py b/src/captionalchemy/tools/audio_analysis/non_speech_detection.py similarity index 94% rename from src/tools/audio_analysis/non_speech_detection.py rename to src/captionalchemy/tools/audio_analysis/non_speech_detection.py index 85f0247..7c353f6 100644 --- a/src/tools/audio_analysis/non_speech_detection.py +++ b/src/captionalchemy/tools/audio_analysis/non_speech_detection.py @@ -1,7 +1,7 @@ from typing import List, Dict import logging -from panns_inference import AudioTagging, labels, SoundEventDetection +from panns_inference import labels, SoundEventDetection import librosa import numpy as np @@ -26,6 +26,9 @@ def detect_non_speech_segments( sed = SoundEventDetection(checkpoint_path=None, device=device) framewise_output = sed.inference(audio) scores = framewise_output[0] + logger.debug( + f"Scores shape: {scores.shape}, Labels: {labels}, Sample rate: {sample_rate}" + ) speech_index = labels.index("Speech") hop_sec = 320 / sample_rate diff --git a/src/tools/audio_analysis/vad.py b/src/captionalchemy/tools/audio_analysis/vad.py similarity index 100% rename from src/tools/audio_analysis/vad.py rename to src/captionalchemy/tools/audio_analysis/vad.py diff --git a/src/tools/captioning/timing_analyzer.py b/src/captionalchemy/tools/captioning/timing_analyzer.py similarity index 97% rename from src/tools/captioning/timing_analyzer.py rename to src/captionalchemy/tools/captioning/timing_analyzer.py index 1f651b2..5f04541 100644 --- a/src/tools/captioning/timing_analyzer.py +++ b/src/captionalchemy/tools/captioning/timing_analyzer.py @@ -14,7 +14,7 @@ class WordTiming: word: str start: float end: float - duration: float = None + duration: Optional[float] = None is_punctuation: bool = False is_sentence_ending: bool = False is_clause_ending: bool = False @@ -169,7 +169,8 @@ def _handle_apostrophes(self, word_timings: List[WordTiming]) -> List[WordTiming # If there is a next token, check if next token is a contraction/possessive suffix if i + 1 < n: next_word = word_timings[i + 1] - # We treat any token whose .word matches eject (apostrophe + s/m/re/ve/ll/d/’t) as a suffix + # Any token whose .word matches eject + # (apostrophe + s/m/re/ve/ll/d/’t) as a suffix if self.contraction_suffix_pattern.match(next_word.word): merged_word = current.word + next_word.word @@ -375,8 +376,8 @@ def _analyze_break_type(self, current_word: str, next_word: str, gap: float) -> def suggest_subtitle_segments( self, word_timings: List[WordTiming], - max_duration: float = None, - max_characters: int = None, + max_duration: Optional[float] = None, + max_characters_per_line: Optional[int] = None, ) -> List[SubtitleSegment]: """ Groups words into subtitle segments based on timing/length. @@ -388,7 +389,9 @@ def suggest_subtitle_segments( word_timings = self._preprocess_whisper_timing(word_timings) max_duration = max_duration or self.max_duration - max_characters = max_characters or (self.max_chars_per_line * self.max_lines) + max_characters = max_characters_per_line or ( + self.max_chars_per_line * self.max_lines + ) analysis = self.calculate_speaking_rate(word_timings) adaptive_max = analysis.get("recommended_max_segment_duration", max_duration) @@ -479,7 +482,8 @@ def _reconstruct_text(self, word_timings: List[WordTiming]) -> str: Does text reconstruction: - Strips each WordTiming.word. - - If it's pure punctuation (one of self.all_punctuation), attach it directly to the previous token. + - If it's pure punctuation (one of self.all_punctuation), + attach it directly to the previous token. - Otherwise, collect it as a normal word. - Finally, join with exactly one space and collapse any accidental extra whitespace. """ diff --git a/src/tools/captioning/transcriber.py b/src/captionalchemy/tools/captioning/transcriber.py similarity index 98% rename from src/tools/captioning/transcriber.py rename to src/captionalchemy/tools/captioning/transcriber.py index 40e001b..044cb54 100644 --- a/src/tools/captioning/transcriber.py +++ b/src/captionalchemy/tools/captioning/transcriber.py @@ -4,7 +4,7 @@ import logging from sys import platform import re -from typing import List, Dict, Optional, Tuple +from typing import List, Optional from dataclasses import dataclass from whisper import Whisper @@ -154,7 +154,7 @@ def transcribe_audio( whisper_build_path: Optional[str] = None, whisper_model_path: Optional[str] = None, device: str = "cuda" if torch.cuda.is_available() else "cpu", - platform: str = platform.lower(), # 'darwin' for macOS, 'linux' for Linux, 'windows' for Windows + platform: str = platform.lower(), # 'darwin', 'linux', 'windows' ) -> List[WordTiming]: """ Transcribe audio to text using OpenAI's Whisper model. diff --git a/src/tools/captioning/words_dictionary.json b/src/captionalchemy/tools/captioning/words_dictionary.json similarity index 100% rename from src/tools/captioning/words_dictionary.json rename to src/captionalchemy/tools/captioning/words_dictionary.json diff --git a/src/captionalchemy/tools/captioning/writers/sami_writer.py b/src/captionalchemy/tools/captioning/writers/sami_writer.py new file mode 100644 index 0000000..7b3c6c0 --- /dev/null +++ b/src/captionalchemy/tools/captioning/writers/sami_writer.py @@ -0,0 +1,212 @@ +import os +import html +from typing import List, Dict, Optional + + +class SAMICaptionWriter: + """ + Accumulates speaker-based transcript segments and writes them out + as a standard .smi (SAMI) caption file following the Microsoft + Synchronized Accessible Media Interchange specification. + """ + + def __init__( + self, + min_silence_duration: float = 1.0, + min_music_duration: float = 1.5, + language: str = "en-US", + title: str = "Generated Captions", + ) -> None: + self._captions: List[Dict] = [] + self.min_silence_duration = min_silence_duration + self.min_music_duration = min_music_duration + self.language = language + self.title = title + self.last_speaker: Optional[str] = None + + def add_caption( + self, + start: float, + end: float, + speaker: Optional[str] = None, + text: Optional[str] = None, + event_type: Optional[str] = None, + label: Optional[str] = None, + ) -> None: + """ + Add one caption entry. + + Args: + start (float): Segment start time in seconds. + end (float): Segment end time in seconds. + speaker (str): Speaker name or ID. + text (str): Transcribed text. + event_type (str, optional): Type of event (speech, music, silence, etc.). + label (str, optional): Original label from audio detection. + """ + # Determine the caption text based on event type + if text: + caption_text = text.strip().replace("\n", " ") + elif event_type == "music": + if end - start >= self.min_music_duration: + caption_text = "[music playing]" + else: + return + elif event_type == "silence": + # Only add silence captions for longer silences + if end - start >= self.min_silence_duration: + caption_text = "[silence]" + else: + return # Skip short silences + elif label: + # For other sound events, use the detected label + caption_text = f"[{label.lower()}]" + else: + caption_text = "[audio]" # Fallback + + self._captions.append( + { + "start": start, + "end": end, + "speaker": speaker, + "text": caption_text, + "event_type": event_type, + } + ) + + def write(self, filepath: str) -> None: + """ + Write all accumulated captions to a .smi file. + + Args: + filepath (str): Path to output .smi file. Overwrites if exists. + """ + # Sort by start time + entries = sorted(self._captions, key=lambda c: c["start"]) + + # Ensure output directory exists + os.makedirs(os.path.dirname(filepath) or ".", exist_ok=True) + + with open(filepath, "w", encoding="utf-8") as f: + # Write SAMI document structure + f.write("\n") + f.write(" \n") + f.write(f" {html.escape(self.title)}\n") + + # Write CSS styles + f.write(' \n") + f.write(" \n") + f.write(" \n") + + # Process entries and add clearing syncs + for i, cap in enumerate(entries): + start_ms = int(cap["start"] * 1000) + end_ms = int(cap["end"] * 1000) + + # Write sync block for the caption + f.write(f' \n') + + # Escape HTML entities in text + escaped_text = html.escape(cap["text"]) + + # Handle line breaks for longer text + if len(escaped_text) > 40: + # Split long text into multiple lines at word boundaries + words = escaped_text.split() + lines = [] + current_line = [] + current_length = 0 + + for word in words: + if current_length + len(word) + 1 <= 40: + current_line.append(word) + current_length += len(word) + 1 + else: + if current_line: + lines.append(" ".join(current_line)) + current_line = [word] + current_length = len(word) + + if current_line: + lines.append(" ".join(current_line)) + + escaped_text = "
".join(lines) + + # Write paragraph with appropriate class and speaker info + if cap["speaker"] and cap["event_type"] == "speech": + # Include speaker information for speech + if self.last_speaker != cap["speaker"]: + self.last_speaker = cap["speaker"] + f.write( + f"

{html.escape(cap['speaker'])}

\n" + ) + f.write(f'

\n') + f.write(f" {escaped_text}\n") + f.write("

\n") + else: + # For music and other events + f.write(f'

\n') + f.write(f" {escaped_text}\n") + f.write("

\n") + + f.write(" \n") + + # Add clearing sync at the end of the caption duration + # Check if there's a gap before the next caption + next_start_ms = None + if i + 1 < len(entries): + next_start_ms = int(entries[i + 1]["start"] * 1000) + + # Add a clearing sync if there's a significant gap or this is the last caption + if ( + next_start_ms is None or next_start_ms - end_ms > 500 + ): # 500ms gap threshold + f.write(f' \n') + f.write(f'

\n') + f.write("  \n") + f.write("

\n") + f.write(" \n") + + f.write(" \n") + f.write("\n") + + def _escape_text(self, text: str) -> str: + """ + Escape HTML entities and handle special characters for SAMI format. + """ + # First escape basic HTML entities + escaped = html.escape(text) + + # Replace other common characters that might cause issues + escaped = escaped.replace('"', """) + escaped = escaped.replace("'", "'") + + return escaped diff --git a/src/tools/captioning/writers/srt_caption_writer.py b/src/captionalchemy/tools/captioning/writers/srt_writer.py similarity index 100% rename from src/tools/captioning/writers/srt_caption_writer.py rename to src/captionalchemy/tools/captioning/writers/srt_writer.py diff --git a/src/captionalchemy/tools/captioning/writers/vtt_writer.py b/src/captionalchemy/tools/captioning/writers/vtt_writer.py new file mode 100644 index 0000000..73dc558 --- /dev/null +++ b/src/captionalchemy/tools/captioning/writers/vtt_writer.py @@ -0,0 +1,118 @@ +import os +from typing import List, Dict, Optional + + +class VTTCaptionWriter: + """ + Accumulates speaker-based transcript segments and writes them out + as a standard .vtt (WebVTT) caption file. + """ + + def __init__( + self, + min_silence_duration: float = 1.0, + min_music_duration: float = 1.5, + ) -> None: + self._captions: List[Dict] = [] + self.min_silence_duration = min_silence_duration + self.min_music_duration = min_music_duration + self.last_speaker: Optional[str] = None + + def add_caption( + self, + start: float, + end: float, + speaker: Optional[str] = None, + text: Optional[str] = None, + event_type: Optional[str] = None, + label: Optional[str] = None, + ) -> None: + """ + Add one caption entry. + + Args: + start (float): Segment start time in seconds. + end (float): Segment end time in seconds. + speaker (str): Speaker name or ID. + text (str): Transcribed text. + event_type (str, optional): Type of event (speech, music, silence, etc.). + label (str, optional): Original label from audio detection. + """ + # Determine the caption text based on event type + if text: + caption_text = text.strip().replace("\n", " ") + elif event_type == "music": + if end - start >= self.min_music_duration: + caption_text = "[MUSIC PLAYING]" + else: + return + elif event_type == "silence": + # Only add silence captions for longer silences + if end - start >= self.min_silence_duration: + caption_text = "[SILENCE]" + else: + return # Skip short silences + elif label: + # For other sound events, use the detected label + caption_text = f"[{label.upper()}]" + else: + caption_text = "[AUDIO]" # Fallback + + self._captions.append( + { + "start": start, + "end": end, + "speaker": speaker, + "text": caption_text, + "event_type": event_type, + } + ) + + def write(self, filepath: str) -> None: + """ + Write all accumulated captions to a .vtt file. + + Args: + filepath (str): Path to output .vtt file. Overwrites if exists. + """ + # sort by start time + entries = sorted(self._captions, key=lambda c: c["start"]) + + # ensure output directory exists + os.makedirs(os.path.dirname(filepath) or ".", exist_ok=True) + + with open(filepath, "w", encoding="utf-8") as f: + # Write WebVTT header + f.write("WEBVTT\n\n") + + for cap in entries: + start_ts = self._format_timestamp(cap["start"]) + end_ts = self._format_timestamp(cap["end"]) + + # Write timing line + f.write(f"{start_ts} --> {end_ts}\n") + + # For speech events, include speaker name using VTT voice tags + if cap["speaker"] and cap["event_type"] == "speech": + if self.last_speaker != cap["speaker"]: + self.last_speaker = cap["speaker"] + f.write(f"{cap['text']}\n\n") + else: + # If the speaker hasn't changed, still use voice tags for consistency + f.write(f"{cap['text']}\n\n") + else: + # For music and other events, just show the description + f.write(f"{cap['text']}\n\n") + + @staticmethod + def _format_timestamp(seconds: float) -> str: + """ + Convert seconds (float) to 'HH:MM:SS.mmm' format for WebVTT. + """ + hours = int(seconds // 3600) + seconds -= hours * 3600 + minutes = int(seconds // 60) + seconds -= minutes * 60 + secs = int(seconds) + millis = int((seconds - secs) * 1000) + return f"{hours:02d}:{minutes:02d}:{secs:02d}.{millis:03d}" diff --git a/src/tools/embed_known_faces.py b/src/captionalchemy/tools/cv/embed_known_faces.py similarity index 80% rename from src/tools/embed_known_faces.py rename to src/captionalchemy/tools/cv/embed_known_faces.py index 440a9db..92d7632 100644 --- a/src/tools/embed_known_faces.py +++ b/src/captionalchemy/tools/cv/embed_known_faces.py @@ -19,7 +19,9 @@ def embed_faces( }, ... ] - This function reads the JSON file, loads the images, and extracts face embeddings using the InsightFace model. + This function reads the JSON file, + loads the images, + and extracts face embeddings using the InsightFace model. Args: known_faces_json (str): Path to the JSON file containing known faces. @@ -27,6 +29,7 @@ def embed_faces( Returns: dict: A dictionary with names as keys and their corresponding face embeddings as values. """ + logger = logging.getLogger(__name__) with open(known_faces_json, "r") as f: known_list = json.load(f) @@ -47,8 +50,6 @@ def embed_faces( allowed_modules=["detection", "recognition"], ) app.prepare(ctx_id=ctx_id, det_size=(640, 640)) - recognizer = app.models["recognition"] - detector = app.models["detection"] embeddings = [] for entry in known_list: @@ -58,12 +59,12 @@ def embed_faces( # Load image img = cv2.imread(image_path) if img is None: - logging.warning(f"Could not read image {image_path}. Skipping.") + logger.warning(f"Could not read image {image_path}. Skipping.") continue faces = app.get(img) if not faces: - logging.warning(f"No faces detected in image {image_path}. Skipping.") + logger.warning(f"No faces detected in image {image_path}. Skipping.") continue embedding = faces[0].embedding @@ -71,4 +72,4 @@ def embed_faces( # Save embeddings to JSON with open(output_embeddings_json, "w") as f: json.dump(embeddings, f, indent=4) - logging.info(f"Embeddings saved to {output_embeddings_json}") + logger.info(f"Embeddings saved to {output_embeddings_json}") diff --git a/src/tools/recognize_faces.py b/src/captionalchemy/tools/cv/recognize_faces.py similarity index 100% rename from src/tools/recognize_faces.py rename to src/captionalchemy/tools/cv/recognize_faces.py diff --git a/src/tools/download_video.py b/src/captionalchemy/tools/media_utils/download_video.py similarity index 99% rename from src/tools/download_video.py rename to src/captionalchemy/tools/media_utils/download_video.py index bdabd8c..f1ead51 100644 --- a/src/tools/download_video.py +++ b/src/captionalchemy/tools/media_utils/download_video.py @@ -1,5 +1,4 @@ import os -import io import uuid import subprocess import tempfile diff --git a/src/tools/extract_audio.py b/src/captionalchemy/tools/media_utils/extract_audio.py similarity index 100% rename from src/tools/extract_audio.py rename to src/captionalchemy/tools/media_utils/extract_audio.py diff --git a/src/tools/audio_analysis/__init__.py b/src/tools/audio_analysis/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/tools/captioning/writers/vtt_caption_writer.py b/src/tools/captioning/writers/vtt_caption_writer.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..58ef86a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,251 @@ +import tempfile +import os +import json +from unittest.mock import Mock, MagicMock +from dataclasses import dataclass +import wave +import math +import logging + +import cv2 +import pytest +import numpy as np + + +from src.captionalchemy.tools.captioning.timing_analyzer import ( + WordTiming, + SubtitleSegment, + TimingAnalyzer, +) +from src.captionalchemy.tools.audio_analysis.audio_segment_integration import ( + AudioEvent, + EventType, +) + + +@pytest.fixture +def temp_dir(): + with tempfile.TemporaryDirectory() as tmpdir: + yield tmpdir + + +@pytest.fixture +def sample_word_timings(): + """ + Sample word timings for testing. + """ + return [ + WordTiming(word="Hello", start=0.0, end=0.5), + WordTiming(word="world", start=0.6, end=1.0), + WordTiming(word="This", start=1.1, end=1.5), + WordTiming(word="is", start=1.6, end=2.0), + WordTiming(word="a", start=2.1, end=2.3), + WordTiming(word="test", start=2.4, end=3.0), + WordTiming( + word=".", start=3.1, end=3.5, is_punctuation=True, is_sentence_ending=True + ), + ] + + +@pytest.fixture +def sample_audio_events(): + """ + Creates sample AudioEvent objects for testing integration. + + This simulates the output from audio analysis components - speech segments, + music detection, and silence periods. This mix helps us test how the + integration logic handles different event types. + """ + return [ + AudioEvent( + start=0.0, end=5.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_00" + ), + AudioEvent(start=5.5, end=8.0, event_type=EventType.MUSIC, label="music"), + AudioEvent( + start=8.5, end=12.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_01" + ), + AudioEvent(start=12.0, end=13.0, event_type=EventType.SILENCE), + AudioEvent( + start=13.5, end=18.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_00" + ), + ] + + +@pytest.fixture +def sample_speech_segments(): + """ + Creates sample speech segments as they would come from VAD. + + Voice Activity Detection (VAD) outputs segments where speech is detected. + This fixture provides realistic timing data for testing. + """ + return [ + {"start": 0.0, "end": 5.0, "duration": 5.0}, + {"start": 8.5, "end": 12.0, "duration": 3.5}, + {"start": 13.5, "end": 18.0, "duration": 4.5}, + ] + + +@pytest.fixture +def sample_diarization_result(): + """ + Creates a sample diarization result. + + This simulates the output from a diarization model, mapping speakers to their + speaking segments. + """ + return { + "SPEAKER_00": {"start": 0.0, "end": 5.0}, + "SPEAKER_01": {"start": 8.5, "end": 12.0}, + "SPEAKER_00": {"start": 13.5, "end": 18.0}, + } + + +@pytest.fixture +def sample_known_faces_json(temp_directory): + """ + Creates a sample known_faces.json file for face recognition testing. + + This simulates the input format for face embedding - a JSON file with + person names and their corresponding image paths. + """ + known_faces_data = [ + {"name": "John Doe", "image_path": os.path.join(temp_directory, "john.jpg")}, + {"name": "Jane Smith", "image_path": os.path.join(temp_directory, "jane.jpg")}, + ] + + json_path = os.path.join(temp_directory, "known_faces.json") + with open(json_path, "w") as f: + json.dump(known_faces_data, f) + + return json_path + + +@pytest.fixture +def mock_whisper_model(): + """ + Mocks the Whisper model for testing transcription. + + This simulates the behavior of the Whisper model without requiring the actual + model files, allowing us to test transcription logic in isolation. + """ + mock_model = MagicMock() + # Mock the transcribe method to return realistic results + mock_result = { + "text": "Hello world. This is a test.", + "segments": [ + { + "words": [ + {"word": " Hello", "start": 0.0, "end": 0.5}, + {"word": " world", "start": 0.6, "end": 1.1}, + {"word": ".", "start": 1.1, "end": 1.2}, + ] + } + ], + } + mock_model.transcribe.return_value = mock_result + return mock_model + + +@pytest.fixture +def mock_face_analysis(): + """ + Mocks the face analysis component for testing speaker identification. + """ + mock_app = MagicMock() + mock_face = MagicMock() + mock_face.embedding = np.random.rand(512) + mock_app.get.return_value = [mock_face] + return mock_app + + +@pytest.fixture +def sample_subtitle_segments(): + """ + Creates sample SubtitleSegment objects for testing. + + These represent the final output of the timing analysis - properly + segmented subtitles with optimal reading speeds and break points. + """ + return [ + SubtitleSegment( + start=0.0, + end=2.5, + text="Hello world.", + word_count=2, + char_count=12, + reading_speed_cps=4.8, + break_reason="sentence_end", + ), + SubtitleSegment( + start=2.5, + end=6.0, + text="This is a test sentence.", + word_count=5, + char_count=24, + reading_speed_cps=6.9, + break_reason="duration_limit", + ), + ] + + +@pytest.fixture +def mock_timing_analyzer(): + """ + Creates a TimingAnalyzer instance with test-friendly settings. + """ + + return TimingAnalyzer( + max_segment_duration=6.0, + min_segment_duration=1.0, + max_reading_speed_cps=6.0, + min_reading_speed_cps=2.0, + ) + + +# Helper functions to create test data + + +def create_test_audio_file(filepath: str, duration_seconds: float = 5.0) -> None: + """ + Creats a minimal .wav audio file for testing. + """ + + sample_rate = 16000 # Standard sample rate + samples = int(sample_rate * duration_seconds) + + with wave.open(filepath, "w") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(sample_rate) + + # Generate a simple sine wave for the audio content + for i in range(samples): + value = int(32767.0 * math.sin(2.0 * math.pi * 440.0 * i / sample_rate)) + wav_file.writeframesraw(value.to_bytes(2, byteorder="little", signed=True)) + + +def create_test_video_file(filepath: str) -> None: + """ + Creates a minimal MP4 video file for testing. + + This function generates a simple video file with a single frame. + """ + + # Create a black image + frame = np.zeros((480, 640, 3), dtype=np.uint8) + + # Write the video file + fourcc = cv2.VideoWriter_fourcc(*"XVID") + out = cv2.VideoWriter(filepath, fourcc, 30.0, (640, 480)) + + out.write(frame) + out.release() + + +@pytest.fixture(autouse=True) +def setup_logging(): + """ + Automatically sets up logging for all tests. + """ + logging.basicConfig(level=logging.WARNING) From c74eec728fd23eec8e439a0af0c56ae4f85c2afa Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Mon, 9 Jun 2025 15:38:34 -0400 Subject: [PATCH 02/21] some tests, need to switch to pytest --- .github/workflows/ci.yml | 80 +++++- .../audio_segment_integration.py | 12 +- .../tools/captioning/timing_analyzer.py | 2 +- .../tools/captioning/transcriber.py | 36 ++- .../tools/captioning/writers/srt_writer.py | 2 +- .../tools/cv/embed_known_faces.py | 2 + tests/conftest.py | 258 ++++++------------ tests/test_audio_integration.py | 141 ++++++++++ tests/test_embed_faces.py | 74 +++++ tests/test_main.py | 172 ++++++++++++ tests/test_srt_writer.py | 71 +++++ tests/test_transcriber.py | 114 ++++++++ 12 files changed, 768 insertions(+), 196 deletions(-) create mode 100644 tests/test_audio_integration.py create mode 100644 tests/test_embed_faces.py create mode 100644 tests/test_main.py create mode 100644 tests/test_srt_writer.py create mode 100644 tests/test_transcriber.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 65fca49..fd6907e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,53 +2,107 @@ name: CI on: push: + branches: + - main + - releases/* tags: - "v*.*.*" + pull_request: + branches: + - main + - releases/* + + workflow_dispatch: + +env: + PYTORCH_VERSION: "2.7.0" jobs: build_and_test: - runs-on: ubuntu-latest + runs-on: ${{ matrix.os }} strategy: matrix: python-version: [3.8, 3.9, 3.10, 3.11] + os: [ubuntu-latest, macos-latest, windows-latest] steps: - name: Checkout code - uses: actions/checkout@v3 + uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies + + - name: Cache pip + uses: actions/cache@v3 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('**/requirements.txt', '**/requirements-dev.txt') }} + restore-keys: | + ${{ runner.os }}-pip-${{ matrix.python-version }}- + ${{ runner.os }}-pip- + + - name: Install dependencies (Ubuntu) + if: runner.os == 'Linux' + run: | + sudo apt-get update + sudo apt-get install -y ffmpeg libsm6 libxext6 libfontconfig1 libxrender1 + # Install audio libraries for pydub + sudo apt-get install -y libavcodec-extra + pip intall torch==${{ env.PYTORCH_VERSION }} torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu + + - name: Install dependencies (macOS) + if: runner.os == 'macOS' + run: | + brew update + brew install ffmpeg + pip3 install torch==${{ env.PYTORCH_VERSION }} torchvision torchaudio + + - name: Install dependencies (Windows) + if: runner.os == 'Windows' + run: | + # Install ffmpeg using chocolatey + choco install -y ffmpeg + # Add ffmpeg to PATH + echo "C:\ProgramData\chocolatey\lib\ffmpeg\tools\ffmpeg\bin" >> $GITHUB_PATH + pip3 install torch==${{ env.PYTORCH_VERSION }} torchvision torchaudio + + - name: Upgrade pip and dependencies run: | - python -m pip install --upgrade pip + python -m pip install --upgrade pip setuptools wheel pip install -r requirements.txt + pip install -r requirements-dev.txt + - name: Lint with flake8 run: | - python -m pip install --upgrade flake8 flake8 src/ tests/ + - name: Run type checks with mypy run: | - python -m pip install --upgrade mypy - mypy --install-types + mypy --install-types --non-interactive --ignore-missing-imports mypy src/ - name: Run tests run: | + pytest --maxfail=1 --disable-warnings -q publish: needs: build_and_test - runs-on: ubuntu-latest if: startsWith(github.ref, 'refs/tags/v') + runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v3 - - name: Set up Python 3.11.5 + uses: actions/checkout@v4 + + - name: Set up Python uses: actions/setup-python@v4 with: - python-version: 3.11.5 - - name: Build distributions + python-version: "3.11" + + - name: Upgrade pip and build run: | - python -m pip install --upgrade build + python -m pip install --upgrade pip build python -m build + - name: Publish to PyPI env: TWINE_USERNAME: __token__ diff --git a/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py b/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py index 8450f8c..eca6308 100644 --- a/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py +++ b/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py @@ -45,6 +45,13 @@ def assign_speakers_to_speech_segment( Returns: List[Dict]: Updated speech segments with assigned speaker IDs. """ + if not diarization: + # If no diarization data, return segments as is + for segment in speech_segments: + segment["speaker_id"] = None + segment["duration"] = segment["end"] - segment["start"] + return speech_segments + speech_segments_with_diarization = [] for segment in speech_segments: @@ -91,9 +98,8 @@ def identify_silence_gaps( """ Identify periods of silence (gaps between speech and non-speech events). - This function finds the "empty spaces" in your audio timeline - periods where - neither speech nor other sounds are detected. These gaps often represent - silence, background noise, or very quiet periods. + This function finds the "empty spaces" in the audio timeline - periods where + neither speech nor other sounds are detected. Args: speech_segments: List of speech segments diff --git a/src/captionalchemy/tools/captioning/timing_analyzer.py b/src/captionalchemy/tools/captioning/timing_analyzer.py index 5f04541..86c0fcd 100644 --- a/src/captionalchemy/tools/captioning/timing_analyzer.py +++ b/src/captionalchemy/tools/captioning/timing_analyzer.py @@ -588,7 +588,7 @@ def _optimize_segments( optimized = [] for seg in segments: # Enforce minimum display time - if seg.duration < self.min_duration: + if seg.duration and seg.duration < self.min_duration: seg.end = seg.start + self.min_duration if optimized: diff --git a/src/captionalchemy/tools/captioning/transcriber.py b/src/captionalchemy/tools/captioning/transcriber.py index 044cb54..2c57153 100644 --- a/src/captionalchemy/tools/captioning/transcriber.py +++ b/src/captionalchemy/tools/captioning/transcriber.py @@ -19,7 +19,7 @@ class WordTiming: word: str start: float end: float - duration: float = None + duration: Optional[float] = None is_punctuation: bool = False is_sentence_ending: bool = False is_clause_ending: bool = False @@ -38,8 +38,18 @@ def __post_init__(self): class Transcriber: """ - Subtitle segmentation that combines speaker identification - with optimal reading timing based on word-level timestamps. + Transcriber class for converting audio to text using OpenAI's Whisper model. + It supports both Whisper Python API and Whisper.cpp for macOS. + It can handle word-level timing and punctuation detection. + It parses lines in the format: + [HH:MM:SS.mmm --> HH:MM:SS.mmm] word + where HH:MM:SS.mmm is the timestamp format. + It returns a list of WordTiming objects containing the word, start time, end time, + duration, and flags for punctuation and subwords. + It also handles subwords by checking the number of leading spaces in the raw token. + It uses ffmpeg to trim audio files to the specified start and end times. + It can run on different platforms (macOS, Linux, Windows) and uses the appropriate + method for transcription based on the platform and availability of Whisper.cpp. """ def __init__(self): @@ -82,6 +92,9 @@ def _parse_timestamps(self, timestamp_str: str) -> float: Returns: float: Total seconds represented by the timestamp. """ + stripped_str = timestamp_str.strip() + if not stripped_str: + raise ValueError("Timestamp string cannot be empty.") try: time_parts = timestamp_str.split(":") hours = int(time_parts[0]) @@ -92,6 +105,17 @@ def _parse_timestamps(self, timestamp_str: str) -> float: seconds = int(seconds_parts[0]) milliseconds = int(seconds_parts[1]) if len(seconds_parts) > 1 else 0 + # Run sanity checks + if hours < 0 or minutes < 0 or seconds < 0 or milliseconds < 0: + raise ValueError("Negative values are not allowed in timestamps.") + if minutes >= 60 or seconds >= 60 or milliseconds >= 1000: + raise ValueError( + "Invalid timestamp values: minutes must be < 60, " + "seconds must be < 60, milliseconds must be < 1000." + ) + if hours > 24: + raise ValueError("Hours must be less than or equal to 24.") + # Calculate total seconds total_seconds = ( hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0 @@ -132,9 +156,8 @@ def _parse_line(self, line: str) -> WordTiming: n_leading_spaces = len(raw_token) - len(raw_token.lstrip(" ")) # Whisper outputs 3 space if whole word, 2 spaces if subword + # Punctuations are taken care of separately is_subword = raw_token not in self.punctuation_set and n_leading_spaces == 2 - if is_subword: - print("Subword detected:", raw_token) token = raw_token.strip() return WordTiming( @@ -143,6 +166,9 @@ def _parse_line(self, line: str) -> WordTiming: end=end_seconds, duration=end_seconds - start_seconds, is_subword=is_subword, + is_punctuation=token in self.punctuation_set, + is_sentence_ending=token in {".", "!", "?"}, + is_clause_ending=token in {",", ";", ":", "—"}, ) def transcribe_audio( diff --git a/src/captionalchemy/tools/captioning/writers/srt_writer.py b/src/captionalchemy/tools/captioning/writers/srt_writer.py index 8fcd13d..19d4060 100644 --- a/src/captionalchemy/tools/captioning/writers/srt_writer.py +++ b/src/captionalchemy/tools/captioning/writers/srt_writer.py @@ -11,7 +11,7 @@ class SRTCaptionWriter: def __init__( self, min_silence_duration: float = 1.0, - min_music_duration: float = 1.5, + min_music_duration: float = 1.0, ) -> None: self._captions: List[Dict] = [] self.min_silence_duration = min_silence_duration diff --git a/src/captionalchemy/tools/cv/embed_known_faces.py b/src/captionalchemy/tools/cv/embed_known_faces.py index 92d7632..d683d78 100644 --- a/src/captionalchemy/tools/cv/embed_known_faces.py +++ b/src/captionalchemy/tools/cv/embed_known_faces.py @@ -30,6 +30,8 @@ def embed_faces( dict: A dictionary with names as keys and their corresponding face embeddings as values. """ logger = logging.getLogger(__name__) + if not known_faces_json: + raise ValueError("known_faces.json must be provided.") with open(known_faces_json, "r") as f: known_list = json.load(f) diff --git a/tests/conftest.py b/tests/conftest.py index 58ef86a..b34c180 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,39 +1,35 @@ -import tempfile import os import json -from unittest.mock import Mock, MagicMock -from dataclasses import dataclass import wave import math -import logging +import tempfile -import cv2 import pytest import numpy as np +import cv2 +from unittest.mock import MagicMock - -from src.captionalchemy.tools.captioning.timing_analyzer import ( - WordTiming, - SubtitleSegment, - TimingAnalyzer, -) -from src.captionalchemy.tools.audio_analysis.audio_segment_integration import ( +from captionalchemy.tools.audio_analysis.audio_segment_integration import ( AudioEvent, EventType, ) +from captionalchemy.tools.captioning.transcriber import WordTiming +from captionalchemy.tools.captioning.timing_analyzer import ( + TimingAnalyzer, + SubtitleSegment, +) @pytest.fixture def temp_dir(): - with tempfile.TemporaryDirectory() as tmpdir: - yield tmpdir + """A temporary directory for files.""" + with tempfile.TemporaryDirectory() as d: + yield d @pytest.fixture def sample_word_timings(): - """ - Sample word timings for testing. - """ + """A short list of WordTiming objects for subtitle-segmentation tests.""" return [ WordTiming(word="Hello", start=0.0, end=0.5), WordTiming(word="world", start=0.6, end=1.0), @@ -47,126 +43,9 @@ def sample_word_timings(): ] -@pytest.fixture -def sample_audio_events(): - """ - Creates sample AudioEvent objects for testing integration. - - This simulates the output from audio analysis components - speech segments, - music detection, and silence periods. This mix helps us test how the - integration logic handles different event types. - """ - return [ - AudioEvent( - start=0.0, end=5.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_00" - ), - AudioEvent(start=5.5, end=8.0, event_type=EventType.MUSIC, label="music"), - AudioEvent( - start=8.5, end=12.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_01" - ), - AudioEvent(start=12.0, end=13.0, event_type=EventType.SILENCE), - AudioEvent( - start=13.5, end=18.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_00" - ), - ] - - -@pytest.fixture -def sample_speech_segments(): - """ - Creates sample speech segments as they would come from VAD. - - Voice Activity Detection (VAD) outputs segments where speech is detected. - This fixture provides realistic timing data for testing. - """ - return [ - {"start": 0.0, "end": 5.0, "duration": 5.0}, - {"start": 8.5, "end": 12.0, "duration": 3.5}, - {"start": 13.5, "end": 18.0, "duration": 4.5}, - ] - - -@pytest.fixture -def sample_diarization_result(): - """ - Creates a sample diarization result. - - This simulates the output from a diarization model, mapping speakers to their - speaking segments. - """ - return { - "SPEAKER_00": {"start": 0.0, "end": 5.0}, - "SPEAKER_01": {"start": 8.5, "end": 12.0}, - "SPEAKER_00": {"start": 13.5, "end": 18.0}, - } - - -@pytest.fixture -def sample_known_faces_json(temp_directory): - """ - Creates a sample known_faces.json file for face recognition testing. - - This simulates the input format for face embedding - a JSON file with - person names and their corresponding image paths. - """ - known_faces_data = [ - {"name": "John Doe", "image_path": os.path.join(temp_directory, "john.jpg")}, - {"name": "Jane Smith", "image_path": os.path.join(temp_directory, "jane.jpg")}, - ] - - json_path = os.path.join(temp_directory, "known_faces.json") - with open(json_path, "w") as f: - json.dump(known_faces_data, f) - - return json_path - - -@pytest.fixture -def mock_whisper_model(): - """ - Mocks the Whisper model for testing transcription. - - This simulates the behavior of the Whisper model without requiring the actual - model files, allowing us to test transcription logic in isolation. - """ - mock_model = MagicMock() - # Mock the transcribe method to return realistic results - mock_result = { - "text": "Hello world. This is a test.", - "segments": [ - { - "words": [ - {"word": " Hello", "start": 0.0, "end": 0.5}, - {"word": " world", "start": 0.6, "end": 1.1}, - {"word": ".", "start": 1.1, "end": 1.2}, - ] - } - ], - } - mock_model.transcribe.return_value = mock_result - return mock_model - - -@pytest.fixture -def mock_face_analysis(): - """ - Mocks the face analysis component for testing speaker identification. - """ - mock_app = MagicMock() - mock_face = MagicMock() - mock_face.embedding = np.random.rand(512) - mock_app.get.return_value = [mock_face] - return mock_app - - @pytest.fixture def sample_subtitle_segments(): - """ - Creates sample SubtitleSegment objects for testing. - - These represent the final output of the timing analysis - properly - segmented subtitles with optimal reading speeds and break points. - """ + """Example SubtitleSegment list from a TimingAnalyzer.""" return [ SubtitleSegment( start=0.0, @@ -191,10 +70,7 @@ def sample_subtitle_segments(): @pytest.fixture def mock_timing_analyzer(): - """ - Creates a TimingAnalyzer instance with test-friendly settings. - """ - + """A TimingAnalyzer configured for testing (no extreme durations).""" return TimingAnalyzer( max_segment_duration=6.0, min_segment_duration=1.0, @@ -203,49 +79,85 @@ def mock_timing_analyzer(): ) -# Helper functions to create test data - +@pytest.fixture +def mock_whisper_model(): + """A fake Whisper model with a stubbed `.transcribe()`.""" + model = MagicMock() + model.transcribe.return_value = { + "text": "Hello world.", + "segments": [ + { + "words": [ + {"word": "Hello", "start": 0.0, "end": 0.5}, + {"word": "world", "start": 0.6, "end": 1.1}, + ] + } + ], + } + return model -def create_test_audio_file(filepath: str, duration_seconds: float = 5.0) -> None: - """ - Creats a minimal .wav audio file for testing. - """ - sample_rate = 16000 # Standard sample rate - samples = int(sample_rate * duration_seconds) +@pytest.fixture +def mock_face_analysis(): + """A fake InsightFace FaceAnalysis instance.""" + app = MagicMock() + face = MagicMock() + face.embedding = np.random.rand(512) + app.get.return_value = [face] + return app - with wave.open(filepath, "w") as wav_file: - wav_file.setnchannels(1) - wav_file.setsampwidth(2) - wav_file.setframerate(sample_rate) - # Generate a simple sine wave for the audio content - for i in range(samples): - value = int(32767.0 * math.sin(2.0 * math.pi * 440.0 * i / sample_rate)) - wav_file.writeframesraw(value.to_bytes(2, byteorder="little", signed=True)) +@pytest.fixture +def sample_known_faces_json(temp_dir): + """Creates a small known_faces.json and empty image files.""" + data = [ + {"name": "Alice", "image_path": os.path.join(temp_dir, "alice.jpg")}, + {"name": "Bob", "image_path": os.path.join(temp_dir, "bob.jpg")}, + ] + # touch the files + for e in data: + open(e["image_path"], "wb").close() + path = os.path.join(temp_dir, "known_faces.json") + with open(path, "w") as f: + json.dump(data, f) + return path -def create_test_video_file(filepath: str) -> None: - """ - Creates a minimal MP4 video file for testing. +@pytest.fixture +def sample_audio_events(): + """A mixed list of AudioEvent for integration logic tests.""" + return [ + AudioEvent( + start=0.0, end=5.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_00" + ), + AudioEvent(start=5.5, end=8.0, event_type=EventType.MUSIC, label="music"), + AudioEvent( + start=8.5, end=12.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_01" + ), + AudioEvent(start=12.0, end=13.0, event_type=EventType.SILENCE), + AudioEvent( + start=13.5, end=18.0, event_type=EventType.SPEECH, speaker_id="SPEAKER_00" + ), + ] - This function generates a simple video file with a single frame. - """ - # Create a black image - frame = np.zeros((480, 640, 3), dtype=np.uint8) +def create_test_audio_file(path: str, duration: float = 2.0): + """Helper to write a tiny WAV sine wave file.""" + sr = 16000 + samples = int(sr * duration) + with wave.open(path, "w") as w: + w.setnchannels(1) + w.setsampwidth(2) + w.setframerate(sr) + for i in range(samples): + val = int(32767 * math.sin(2 * math.pi * 440 * i / sr)) + w.writeframesraw(val.to_bytes(2, "little", signed=True)) - # Write the video file - fourcc = cv2.VideoWriter_fourcc(*"XVID") - out = cv2.VideoWriter(filepath, fourcc, 30.0, (640, 480)) +def create_test_video_file(path: str): + """Helper to write a one-frame MP4 via OpenCV.""" + frame = np.zeros((480, 640, 3), dtype=np.uint8) + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + out = cv2.VideoWriter(path, fourcc, 24.0, (640, 480)) out.write(frame) out.release() - - -@pytest.fixture(autouse=True) -def setup_logging(): - """ - Automatically sets up logging for all tests. - """ - logging.basicConfig(level=logging.WARNING) diff --git a/tests/test_audio_integration.py b/tests/test_audio_integration.py new file mode 100644 index 0000000..9068a4c --- /dev/null +++ b/tests/test_audio_integration.py @@ -0,0 +1,141 @@ +import unittest + +from captionalchemy.tools.audio_analysis.audio_segment_integration import ( + AudioEvent, + EventType, + assign_speakers_to_speech_segment, + identify_silence_gaps, +) + + +class TestAudioEvent(unittest.TestCase): + def test_audio_event_creation_speech(self): + event = AudioEvent( + start=0.0, + end=1.0, + speaker_id="Speaker_00", + event_type=EventType.SPEECH, + confidence=0.95, + ) + self.assertEqual(event.start, 0.0) + self.assertEqual(event.end, 1.0) + self.assertEqual(event.speaker_id, "Speaker_00") + self.assertEqual(event.event_type, EventType.SPEECH) + self.assertEqual(event.confidence, 0.95) + + def test_audio_event_creation_music(self): + """Test creating a music AudioEvent.""" + event = AudioEvent( + start=10.0, + end=15.0, + event_type=EventType.MUSIC, + label="background_music", + confidence=0.88, + ) + + self.assertEqual(event.event_type, EventType.MUSIC) + self.assertEqual(event.label, "background_music") + self.assertIsNone(event.speaker_id) + + def test_audio_event_duration_auto_calculation(self): + """Test that duration is automatically calculated.""" + event = AudioEvent(start=1.5, end=4.7, event_type=EventType.SILENCE) + self.assertAlmostEqual(event.duration, 3.2, places=6) + + def test_audio_event_validation_invalid_timing(self): + """Test that invalid timing raises ValueError.""" + with self.assertRaises(ValueError) as context: + AudioEvent(start=2.0, end=1.0, event_type=EventType.SPEECH) + + self.assertIn( + "End time must be greater than start time", str(context.exception) + ) + + +class TestSpeakerAssignment(unittest.TestCase): + def test_assign_speakers(self): + speech_segments = [{"start": 0.0, "end": 5.0}, {"start": 10.0, "end": 15.0}] + + diarization = { + "SPEAKER_00": {"start": 0.0, "end": 5.0}, + "SPEAKER_01": {"start": 10.0, "end": 15.0}, + } + result = assign_speakers_to_speech_segment(speech_segments, diarization) + + self.assertEqual(len(result), 2) + self.assertEqual(result[0]["speaker_id"], "SPEAKER_00") + self.assertEqual(result[1]["speaker_id"], "SPEAKER_01") + + def test_assign_speakers_overlapping_segments(self): + """Test speaker assignment with overlapping segments.""" + speech_segments = [{"start": 2.0, "end": 7.0}] # Overlaps both speakers + + diarization = { + "SPEAKER_00": {"start": 0.0, "end": 5.0}, # 3 seconds overlap + "SPEAKER_01": {"start": 4.0, "end": 10.0}, # 3 seconds overlap + } + + result = assign_speakers_to_speech_segment(speech_segments, diarization) + self.assertEqual(len(result), 1) + self.assertEqual( + result[0]["speaker_id"], "SPEAKER_00" + ) # Overlapping segments should assign the first speaker + + def test_assign_speakers_no_diarization(self): + """Test behavior when no diarization data is provided.""" + speech_segments = [{"start": 0.0, "end": 5.0}] + diarization = {} + + result = assign_speakers_to_speech_segment(speech_segments, diarization) + + self.assertEqual(len(result), 1) + self.assertIsNone(result[0].get("speaker_id")) + + def test_assign_speakers_duration_calculation(self): + """Test that duration is calculated correctly.""" + speech_segments = [{"start": 1.5, "end": 4.7}] + diarization = {"SPEAKER_00": {"start": 0.0, "end": 10.0}} + result = assign_speakers_to_speech_segment(speech_segments, diarization) + self.assertAlmostEqual(result[0]["duration"], 3.2, places=6) + + +class TestSilenceIdentification(unittest.TestCase): + def test_identify_silence_gaps(self): + speech_segments = [ + {"start": 0.0, "end": 3.0}, + {"start": 5.0, "end": 8.0}, # 2-second gap + ] + non_speech_segments = [] + silences = identify_silence_gaps( + speech_segments, + non_speech_segments, + total_audio_duration=10.0, + min_silence_duration=0.5, + ) + self.assertEqual(len(silences), 2) + + first, second = silences + self.assertAlmostEqual(first["start"], 3.0) + self.assertAlmostEqual(first["end"], 5.0) + self.assertAlmostEqual(first["duration"], 2.0) + self.assertAlmostEqual(first["event_type"], EventType.SILENCE) + + self.assertAlmostEqual(second["start"], 8.0) + self.assertAlmostEqual(second["end"], 10.0) + self.assertAlmostEqual(second["duration"], 2.0) + self.assertAlmostEqual(second["event_type"], EventType.SILENCE) + + def test_no_silence_when_too_short(self): + speech_segments = [{"start": 0.0, "end": 1.0}, {"start": 1.4, "end": 2.0}] + non_speech_segments = [] + silences = identify_silence_gaps( + speech_segments, + non_speech_segments, + total_audio_duration=3.0, + min_silence_duration=0.5, + ) + self.assertEqual(len(silences), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_embed_faces.py b/tests/test_embed_faces.py new file mode 100644 index 0000000..66f7d6e --- /dev/null +++ b/tests/test_embed_faces.py @@ -0,0 +1,74 @@ +import unittest +import tempfile +import os +import json +import numpy as np +from unittest.mock import patch, MagicMock + +from captionalchemy.tools.cv.embed_known_faces import embed_faces + + +class TestEmbedFacesFunctionality(unittest.TestCase): + def setUp(self): + # Two dummy face entries + self.entries = [ + {"name": "Alice", "image_path": "path/to/alice.jpg"}, + {"name": "Bob", "image_path": "path/to/bob.jpg"}, + ] + + def test_embed_faces_empty_json(self): + with self.assertRaises(ValueError): + embed_faces("", "output_embeddings.json") + + @patch("captionalchemy.tools.cv.embed_known_faces.FaceAnalysis") + @patch("captionalchemy.tools.cv.embed_known_faces.cv2.imread") + @patch("torch.cuda.is_available", return_value=False) + @patch("torch.backends.mps.is_available", return_value=False) + def test_embed_faces_success( + self, mock_mps, mock_cuda, mock_imread, mock_face_analysis + ): + # Prepare temporary JSON input and output paths + with tempfile.TemporaryDirectory() as tmpdir: + input_path = os.path.join(tmpdir, "known_faces.json") + output_path = os.path.join(tmpdir, "output_embeddings.json") + + # Write sample known-faces JSON + for entry in self.entries: + # Create dummy image files + img_path = os.path.join(tmpdir, f"{entry['name']}.jpg") + open(img_path, "wb").close() + entry["image_path"] = img_path + with open(input_path, "w") as f: + json.dump(self.entries, f) + + # Mock cv2.imread to return a dummy array for any path + mock_imread.return_value = np.random.randint( + 0, 255, (100, 100, 3), dtype=np.uint8 + ) + + # Mock FaceAnalysis to return a dummy embedding and get() method + mock_app = MagicMock() + # Each call returns a single face with a known embedding + face = MagicMock() + face.embedding = np.random.rand(512).astype(np.float32) + mock_app.get.return_value = [face] + mock_face_analysis.return_value = mock_app + + # Run + embed_faces(input_path, output_path) + + # Read output and verify + with open(output_path, "r") as f: + data = json.load(f) + # Should have two entries + self.assertEqual(len(data), 2) + names = [entry["name"] for entry in data] + self.assertCountEqual(names, ["Alice", "Bob"]) + for d in data: + emb = d["embedding"] + self.assertIsInstance(emb, list) + self.assertEqual(len(emb), 512) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..c706c73 --- /dev/null +++ b/tests/test_main.py @@ -0,0 +1,172 @@ +import unittest +import tempfile +import os +import json +import uuid +from unittest.mock import Mock, patch, MagicMock, call +from typing import List, Dict, Any +import numpy as np +import shutil + + +class TestMainPipelineIntegration(unittest.TestCase): + def setUp(self): + self.test_video_url = "test_video.mp4" + self.temp_dir = tempfile.mkdtemp() + self.setup_mocks() + + def tearDown(self): + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def setup_mocks(self): + """Set up mocks for pipeline testing.""" + self.mock_video_manager = Mock() + self.mock_video_manager.get_video_data.return_value = ( + "/path/to/downloaded/video.mp4" + ) + + self.mock_extract_audio = Mock() + + # Mock speech segments (VAD) + self.mock_speech_segments = [ + {"start": 0.0, "end": 5.0, "duration": 5.0}, + {"start": 10.0, "end": 15.0, "duration": 5.0}, + {"start": 20.0, "end": 25.0, "duration": 5.0}, + ] + + # Mock diarization result + self.mock_diarization_result = { + "SPEAKER_00": {"start": 0.0, "end": 15.0}, + "SPEAKER_01": {"start": 10.0, "end": 25.0}, + } + + # Mock non-speech events + self.mock_non_speech_events = [ + { + "start": 5.5, + "end": 8.0, + "label": "music", + "confidence": 0.9, + "duration": 2.5, + } + ] + + # Mock face recognition results + self.mock_recognized_faces = [ + { + "timesatmp": 0.0, + "bbox": [100, 100, 200, 200], + "face_id": "face_1", + "name": "John Doe", + }, + { + "timesatmp": 10.0, + "bbox": [150, 150, 250, 250], + "face_id": "face_2", + "name": "Jane Smith", + }, + ] + + # Mock word timings from transcriber + self.mock_word_timings = [ + {"word": "Hello", "start": 0.0, "end": 0.5, "duration": 0.5}, + {"word": "world", "start": 0.6, "end": 1.1, "duration": 0.5}, + {"word": ".", "start": 1.1, "end": 1.2, "duration": 0.1}, + {"word": "This", "start": 1.3, "end": 1.8, "duration": 0.5}, + {"word": "is", "start": 1.9, "end": 2.0, "duration": 0.1}, + {"word": "a", "start": 2.1, "end": 2.2, "duration": 0.1}, + {"word": "test", "start": 2.3, "end": 2.8, "duration": 0.5}, + ] + + # Mock subtitle segments from timing analyzer + self.mock_subtitle_segments = [ + { + "start": 0.0, + "end": 1.2, + "text": "Hello world.", + "word_count": 2, + "char_count": 12, + "break_reason": "sentence_ending", + } + ] + + +@patch("tools.captioning.timing_analyzer.TimingAnalyzer") +@patch("tools.captioning.transcriber.Transcriber") +@patch("src.captionalchemy.tools.cv.recognize_faces.recognize_faces") +@patch("tools.audio_analysis.audio_segment_integration.integrate_audio_segments") +@patch("tools.audio_analysis.non_speech_detection.detect_non_speech_segments") +@patch("tools.audio_analysis.diarization.diarize") +@patch("tools.audio_analysis.vad.get_speech_segments") +@patch("src.captionalchemy.tools.media_utils.extract_audio.extract_audio") +@patch("src.captionalchemy.tools.media_utils.download_video.VideoManager") +@patch("src.captionalchemy.tools.cv.embed_known_faces.embed_faces") +@patch("whisper.load_model") +@patch("torch.cuda.is_available") +def test_complete_pipeline_success( + self, + mock_cuda, + mock_whisper_load, + mock_embed_faces, + mock_video_manager_class, + mock_extract_audio, + mock_vad, + mock_diarize, + mock_non_speech, + mock_integrate_audio, + mock_recognize_faces, + mock_transcriber_class, + mock_timing_analyzer_class, + mock_srt_writer_class, +): + """Test the complete pipeline with successful execution.""" + mock_cuda.return_value = True + + mock_whisper_model = Mock() + mock_whisper_load.return_value = mock_whisper_model + + mock_embed_faces.return_value = None + + mock_video_manager = Mock() + mock_video_manager.get_video_data.return_value = "/tmp/video.mp4" + mock_video_manager_class.return_value = mock_video_manager + + mock_extract_audio.return_value = None + + mock_vad.return_value = self.mock_speech_segments + mock_diarize.return_value = self.mock_diarization_result + mock_non_speech.return_value = self.mock_non_speech_events + + mock_integrated_events = [] + + for segment in self.mock_speech_segments: + mock_event = Mock() + mock_event.event_type.value = "speech" + mock_event.start = segment["start"] + mock_event.end = segment["end"] + mock_event.speaker_id = "SPEAKER_00" if segment["start"] < 15 else "SPEAKER_01" + mock_integrated_events.append(mock_event) + + mock_integrate_audio.return_value = mock_integrated_events + + # Set up mock face recognition + mock_recognize_faces.return_value = self.mock_recognized + + mock_transcriber = Mock() + mock_transcriber.transcribe_audio.return_value = [ + Mock(**wt) for wt in self.mock_word_timings + ] + mock_transcriber_class.return_value = mock_transcriber + + mock_timing_analyzer = Mock() + mock_timing_analyzer.suggest_subtitle_segments.return_value = [ + Mock(**seg) for seg in self.mock_subtitle_segments + ] + mock_timing_analyzer_class.return_value = mock_timing_analyzer + + # Mock SRT writer + mock_srt_writer = Mock() + mock_srt_writer_class.return_value = mock_srt_writer + + # Create simplified mock of main function diff --git a/tests/test_srt_writer.py b/tests/test_srt_writer.py new file mode 100644 index 0000000..e083e70 --- /dev/null +++ b/tests/test_srt_writer.py @@ -0,0 +1,71 @@ +import unittest +import tempfile +import os + +from captionalchemy.tools.captioning.writers.srt_writer import SRTCaptionWriter + + +class TestSRTWriter(unittest.TestCase): + def test_add_caption_text_and_silence_music(self): + writer = SRTCaptionWriter() + + # Speech with text + writer.add_caption( + 0.0, 1.0, speaker="SPEAKER_00", text="Hello world.", event_type="speech" + ) + # Short silence (should skip) + writer.add_caption( + 1.0, 1.2, speaker="SPEAKER_00", text="", event_type="silence" + ) + # Long silence (should add) + writer.add_caption( + 1.2, 3.0, speaker="SPEAKER_00", text="", event_type="silence" + ) + # Music event (should skip) + writer.add_caption(3.0, 4.5, event_type="music") + # Long music event (should add) + writer.add_caption(4.5, 6.0, event_type="music") + + caps = writer._captions + # Expect 1 speech, 1 silence, and 1 music caption, 3 entries total + self.assertEqual(len(caps), 3) + self.assertEqual(caps[0].text, "Hello world.") + self.assertEqual(caps[1]["text"], "[SILENCE]") + self.assertEqual(caps[2]["text"], "[MUSIC PLAYING]") + + def test_write_srt_file(self): + writer = SRTCaptionWriter() + + # Two entires with same speaker + writer.add_caption( + 0.0, 1.0, speaker="SPEAKER_00", text="Hello world.", event_type="speech" + ) + writer.add_caption( + 1.0, 2.0, speaker="SPEAKER_00", text="This is a test.", event_type="speech" + ) + # Different speaker + writer.add_caption( + 2.0, 3.0, speaker="SPEAKER_01", text="Goodbye.", event_type="speech" + ) + + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".srt") + tmp.close() + path = tmp.name + + try: + writer.write(path) + content = open(path, "r", encoding="utf-8").read().splitlines() + # Check index and timings in first block + self.assertEqual(content[0], "1") + self.assertTrue("-->" in content[1]) + # Speaker line for first caption + self.assertTrue("SPEAKER_00: Hello world." in content[2]) + # Next blocks exist + self.assertIn("2", content) + self.assertIn("3", content) + finally: + os.remove(path) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_transcriber.py b/tests/test_transcriber.py new file mode 100644 index 0000000..2c67a0d --- /dev/null +++ b/tests/test_transcriber.py @@ -0,0 +1,114 @@ +import unittest +import tempfile +import os +import subprocess +from unittest.mock import patch, Mock +import numpy as np +import torch + +from captionalchemy.tools.captioning.transcriber import Transcriber, WordTiming + + +class TestWordTimingDataclass(unittest.TestCase): + def test_word_timing_creation(self): + """Test creating a WordTiming instance.""" + word_timing = WordTiming(word="test", start=0.0, end=1.0, is_subword=True) + self.assertEqual(word_timing.word, "test") + self.assertEqual(word_timing.start, 0.0) + self.assertEqual(word_timing.end, 1.0) + self.assertFalse(word_timing.is_punctuation) + self.assertFalse(word_timing.is_sentence_ending) + self.assertTrue(word_timing.is_subword) + + def test_word_timing_invalid(self): + with self.assertRaises(ValueError): + WordTiming(word="", start=0.0, end=1.0) # empty word + with self.assertRaises(ValueError): + WordTiming(word="hi", start=-1.0, end=0.5) # negative start + with self.assertRaises(ValueError): + WordTiming(word="hi", start=1.0, end=0.5) # end < start + + +class TestTranscriberParse(unittest.TestCase): + def setUp(self): + self.transcriber = Transcriber() + + def test_parse_timestamps(self): + """Test parsing timestamps from a Whisper cpp output.""" + self.assertAlmostEqual( + self.transcriber._parse_timestamps("00:01:23.456"), 83.456, places=3 + ) + self.assertAlmostEqual( + self.transcriber._parse_timestamps("2:03:04.007"), + 2 * 3600 + 3 * 60 + 4.007, + places=3, + ) + with self.assertRaises(ValueError): + self.transcriber._parse_timestamps("99:99:99") + + def test_parse_line(self): + """Test parsing line from Whisper cpp output.""" + line = "[00:00:00.000 --> 00:00:00.500] Hello" + wt = self.transcriber._parse_line(line) + self.assertEqual(wt.word, "Hello") + self.assertAlmostEqual(wt.start, 0.0, places=3) + self.assertAlmostEqual(wt.end, 0.5, places=3) + # punctuation + line2 = "[00:00:01.000 --> 00:00:01.200] ." + wt2 = self.transcriber._parse_line(line2) + self.assertEqual(wt2.word, ".") + self.assertTrue(wt2.is_punctuation) + + @patch("tempfile.NamedTemporaryFile") + @patch("subprocess.run") + @patch("whisper.load_model") + @patch("os.remove") + def test_transcribe_audio_python_api( + self, mock_remove, mock_load, mock_run, mock_tempfile + ): + """Test transcribing audio using Whisper Python API.""" + # Mock temp file + fake = Mock() + fake.name = "/tmp/fake.wav" + mock_tempfile.return_value = fake + + # Mock ffmpeg trim + mock_run.return_value = Mock(returncode=0) + + # Mock whisper model load & transcribe + mock_model = Mock() + mock_load.return_value = mock_model + mock_model.transcribe.return_value = { + "text": "Hello world.", + "segments": [ + { + "words": [ + {"word": "Hello", "start": 0.0, "end": 0.5}, + {"word": "world", "start": 0.6, "end": 1.1}, + ] + } + ], + } + + transcriber = Transcriber() + words = transcriber.transcribe_audio( + audio_file="in.wav", + start=1.0, + end=2.0, + model="base", + whisper_build_path=None, + whisper_model_path=None, + platform="linux", + ) + + self.assertEqual(len(words), 2) + self.assertEqual(words[0].word, "Hello") + self.assertAlmostEqual(words[0].start, 0.0, places=3) + self.assertAlmostEqual(words[0].end, 0.5, places=3) + self.assertEqual(words[1].word, "world") + self.assertAlmostEqual(words[1].start, 0.6, places=3) + self.assertAlmostEqual(words[1].end, 1.1, places=3) + + +if __name__ == "__main__": + unittest.main() From cc82fd539dbcc54fc02eb4c4dd3cba5c24ed3889 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Mon, 9 Jun 2025 21:47:31 -0400 Subject: [PATCH 03/21] wrapping up tests --- .github/workflows/ci.yml | 2 +- tests/conftest.py | 81 ++- tests/test_audio_integration.py | 260 +++++---- tests/test_embed_faces.py | 110 ++-- tests/test_main.py | 924 +++++++++++++++++++++++++++----- tests/test_srt_writer.py | 107 ++-- tests/test_transcriber.py | 212 ++++---- 7 files changed, 1201 insertions(+), 495 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fd6907e..bae5806 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -83,7 +83,7 @@ jobs: mypy src/ - name: Run tests run: | - pytest --maxfail=1 --disable-warnings -q + pytest --cov=captionalchemy -v publish: needs: build_and_test diff --git a/tests/conftest.py b/tests/conftest.py index b34c180..35a0a8d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -81,7 +81,7 @@ def mock_timing_analyzer(): @pytest.fixture def mock_whisper_model(): - """A fake Whisper model with a stubbed `.transcribe()`.""" + """A fake Whisper model with a stubbed .transcribe().""" model = MagicMock() model.transcribe.return_value = { "text": "Hello world.", @@ -141,6 +141,22 @@ def sample_audio_events(): ] +@pytest.fixture +def sample_test_audio_file(temp_dir): + """Creates a test WAV file and returns the path.""" + path = os.path.join(temp_dir, "test_audio.wav") + create_test_audio_file(path) + return path + + +@pytest.fixture +def sample_test_video_file(temp_dir): + """Creates a test MP4 file and returns the path.""" + path = os.path.join(temp_dir, "test_video.mp4") + create_test_video_file(path) + return path + + def create_test_audio_file(path: str, duration: float = 2.0): """Helper to write a tiny WAV sine wave file.""" sr = 16000 @@ -161,3 +177,66 @@ def create_test_video_file(path: str): out = cv2.VideoWriter(path, fourcc, 24.0, (640, 480)) out.write(frame) out.release() + + +# Additional fixtures for pytest conversion + + +@pytest.fixture +def mock_speech_segments(): + """Mock speech segments for testing.""" + return [ + {"start": 0.0, "end": 5.0, "duration": 5.0}, + {"start": 10.0, "end": 15.0, "duration": 5.0}, + {"start": 20.0, "end": 25.0, "duration": 5.0}, + ] + + +@pytest.fixture +def mock_diarization_result(): + """Mock diarization result.""" + return { + "SPEAKER_00": {"start": 0.0, "end": 15.0}, + "SPEAKER_01": {"start": 10.0, "end": 25.0}, + } + + +@pytest.fixture +def mock_face_recognition_results(): + """Mock face recognition results.""" + return [ + { + "timestamp": 0.0, + "bbox": [100, 100, 200, 200], + "face_id": "face_1", + "name": "John Doe", + }, + { + "timestamp": 10.0, + "bbox": [150, 150, 250, 250], + "face_id": "face_2", + "name": "Jane Smith", + }, + ] + + +@pytest.fixture +def mock_non_speech_events(): + """Mock non-speech audio events.""" + return [ + { + "start": 5.5, + "end": 8.0, + "label": "music", + "confidence": 0.9, + "duration": 2.5, + } + ] + + +# Pytest configuration for better test output +def pytest_configure(config): + """Configure pytest with custom settings.""" + # Add custom markers if needed + config.addinivalue_line("markers", "integration: mark test as integration test") + config.addinivalue_line("markers", "slow: mark test as slow running") diff --git a/tests/test_audio_integration.py b/tests/test_audio_integration.py index 9068a4c..95ce350 100644 --- a/tests/test_audio_integration.py +++ b/tests/test_audio_integration.py @@ -1,4 +1,4 @@ -import unittest +import pytest from captionalchemy.tools.audio_analysis.audio_segment_integration import ( AudioEvent, @@ -8,134 +8,130 @@ ) -class TestAudioEvent(unittest.TestCase): - def test_audio_event_creation_speech(self): - event = AudioEvent( - start=0.0, - end=1.0, - speaker_id="Speaker_00", - event_type=EventType.SPEECH, - confidence=0.95, - ) - self.assertEqual(event.start, 0.0) - self.assertEqual(event.end, 1.0) - self.assertEqual(event.speaker_id, "Speaker_00") - self.assertEqual(event.event_type, EventType.SPEECH) - self.assertEqual(event.confidence, 0.95) - - def test_audio_event_creation_music(self): - """Test creating a music AudioEvent.""" - event = AudioEvent( - start=10.0, - end=15.0, - event_type=EventType.MUSIC, - label="background_music", - confidence=0.88, - ) - - self.assertEqual(event.event_type, EventType.MUSIC) - self.assertEqual(event.label, "background_music") - self.assertIsNone(event.speaker_id) - - def test_audio_event_duration_auto_calculation(self): - """Test that duration is automatically calculated.""" - event = AudioEvent(start=1.5, end=4.7, event_type=EventType.SILENCE) - self.assertAlmostEqual(event.duration, 3.2, places=6) - - def test_audio_event_validation_invalid_timing(self): - """Test that invalid timing raises ValueError.""" - with self.assertRaises(ValueError) as context: - AudioEvent(start=2.0, end=1.0, event_type=EventType.SPEECH) - - self.assertIn( - "End time must be greater than start time", str(context.exception) - ) - - -class TestSpeakerAssignment(unittest.TestCase): - def test_assign_speakers(self): - speech_segments = [{"start": 0.0, "end": 5.0}, {"start": 10.0, "end": 15.0}] - - diarization = { - "SPEAKER_00": {"start": 0.0, "end": 5.0}, - "SPEAKER_01": {"start": 10.0, "end": 15.0}, - } - result = assign_speakers_to_speech_segment(speech_segments, diarization) - - self.assertEqual(len(result), 2) - self.assertEqual(result[0]["speaker_id"], "SPEAKER_00") - self.assertEqual(result[1]["speaker_id"], "SPEAKER_01") - - def test_assign_speakers_overlapping_segments(self): - """Test speaker assignment with overlapping segments.""" - speech_segments = [{"start": 2.0, "end": 7.0}] # Overlaps both speakers - - diarization = { - "SPEAKER_00": {"start": 0.0, "end": 5.0}, # 3 seconds overlap - "SPEAKER_01": {"start": 4.0, "end": 10.0}, # 3 seconds overlap - } - - result = assign_speakers_to_speech_segment(speech_segments, diarization) - self.assertEqual(len(result), 1) - self.assertEqual( - result[0]["speaker_id"], "SPEAKER_00" - ) # Overlapping segments should assign the first speaker - - def test_assign_speakers_no_diarization(self): - """Test behavior when no diarization data is provided.""" - speech_segments = [{"start": 0.0, "end": 5.0}] - diarization = {} - - result = assign_speakers_to_speech_segment(speech_segments, diarization) - - self.assertEqual(len(result), 1) - self.assertIsNone(result[0].get("speaker_id")) - - def test_assign_speakers_duration_calculation(self): - """Test that duration is calculated correctly.""" - speech_segments = [{"start": 1.5, "end": 4.7}] - diarization = {"SPEAKER_00": {"start": 0.0, "end": 10.0}} - result = assign_speakers_to_speech_segment(speech_segments, diarization) - self.assertAlmostEqual(result[0]["duration"], 3.2, places=6) - - -class TestSilenceIdentification(unittest.TestCase): - def test_identify_silence_gaps(self): - speech_segments = [ - {"start": 0.0, "end": 3.0}, - {"start": 5.0, "end": 8.0}, # 2-second gap - ] - non_speech_segments = [] - silences = identify_silence_gaps( - speech_segments, - non_speech_segments, - total_audio_duration=10.0, - min_silence_duration=0.5, - ) - self.assertEqual(len(silences), 2) - - first, second = silences - self.assertAlmostEqual(first["start"], 3.0) - self.assertAlmostEqual(first["end"], 5.0) - self.assertAlmostEqual(first["duration"], 2.0) - self.assertAlmostEqual(first["event_type"], EventType.SILENCE) - - self.assertAlmostEqual(second["start"], 8.0) - self.assertAlmostEqual(second["end"], 10.0) - self.assertAlmostEqual(second["duration"], 2.0) - self.assertAlmostEqual(second["event_type"], EventType.SILENCE) - - def test_no_silence_when_too_short(self): - speech_segments = [{"start": 0.0, "end": 1.0}, {"start": 1.4, "end": 2.0}] - non_speech_segments = [] - silences = identify_silence_gaps( - speech_segments, - non_speech_segments, - total_audio_duration=3.0, - min_silence_duration=0.5, - ) - self.assertEqual(len(silences), 0) - - -if __name__ == "__main__": - unittest.main() +def test_audio_event_creation_speech(): + event = AudioEvent( + start=0.0, + end=1.0, + speaker_id="Speaker_00", + event_type=EventType.SPEECH, + confidence=0.95, + ) + assert event.start == 0.0 + assert event.end == 1.0 + assert event.speaker_id == "Speaker_00" + assert event.event_type == EventType.SPEECH + assert event.confidence == 0.95 + + +def test_audio_event_creation_music(): + """Test creating a music AudioEvent.""" + event = AudioEvent( + start=10.0, + end=15.0, + event_type=EventType.MUSIC, + label="background_music", + confidence=0.88, + ) + + assert event.event_type == EventType.MUSIC + assert event.label == "background_music" + assert event.speaker_id is None + + +def test_audio_event_duration_auto_calculation(): + """Test that duration is automatically calculated.""" + event = AudioEvent(start=1.5, end=4.7, event_type=EventType.SILENCE) + assert abs(event.duration - 3.2) < 1e-6 + + +def test_audio_event_validation_invalid_timing(): + """Test that invalid timing raises ValueError.""" + with pytest.raises(ValueError, match="End time must be greater than start time"): + AudioEvent(start=2.0, end=1.0, event_type=EventType.SPEECH) + + +def test_assign_speakers(): + speech_segments = [{"start": 0.0, "end": 5.0}, {"start": 10.0, "end": 15.0}] + + diarization = { + "SPEAKER_00": {"start": 0.0, "end": 5.0}, + "SPEAKER_01": {"start": 10.0, "end": 15.0}, + } + result = assign_speakers_to_speech_segment(speech_segments, diarization) + + assert len(result) == 2 + assert result[0]["speaker_id"] == "SPEAKER_00" + assert result[1]["speaker_id"] == "SPEAKER_01" + + +def test_assign_speakers_overlapping_segments(): + """Test speaker assignment with overlapping segments.""" + speech_segments = [{"start": 2.0, "end": 7.0}] # Overlaps both speakers + + diarization = { + "SPEAKER_00": {"start": 0.0, "end": 5.0}, # 3 seconds overlap + "SPEAKER_01": {"start": 4.0, "end": 10.0}, # 3 seconds overlap + } + + result = assign_speakers_to_speech_segment(speech_segments, diarization) + assert len(result) == 1 + assert ( + result[0]["speaker_id"] == "SPEAKER_00" + ) # Overlapping segments should assign the first speaker + + +def test_assign_speakers_no_diarization(): + """Test behavior when no diarization data is provided.""" + speech_segments = [{"start": 0.0, "end": 5.0}] + diarization = {} + + result = assign_speakers_to_speech_segment(speech_segments, diarization) + + assert len(result) == 1 + assert result[0].get("speaker_id") is None + + +def test_assign_speakers_duration_calculation(): + """Test that duration is calculated correctly.""" + speech_segments = [{"start": 1.5, "end": 4.7}] + diarization = {"SPEAKER_00": {"start": 0.0, "end": 10.0}} + result = assign_speakers_to_speech_segment(speech_segments, diarization) + assert abs(result[0]["duration"] - 3.2) < 1e-6 + + +def test_identify_silence_gaps(): + speech_segments = [ + {"start": 0.0, "end": 3.0}, + {"start": 5.0, "end": 8.0}, # 2-second gap + ] + non_speech_segments = [] + silences = identify_silence_gaps( + speech_segments, + non_speech_segments, + total_audio_duration=10.0, + min_silence_duration=0.5, + ) + assert len(silences) == 2 + + first, second = silences + assert abs(first["start"] - 3.0) < 1e-6 + assert abs(first["end"] - 5.0) < 1e-6 + assert abs(first["duration"] - 2.0) < 1e-6 + assert first["event_type"] == EventType.SILENCE + + assert abs(second["start"] - 8.0) < 1e-6 + assert abs(second["end"] - 10.0) < 1e-6 + assert abs(second["duration"] - 2.0) < 1e-6 + assert second["event_type"] == EventType.SILENCE + + +def test_no_silence_when_too_short(): + speech_segments = [{"start": 0.0, "end": 1.0}, {"start": 1.4, "end": 2.0}] + non_speech_segments = [] + silences = identify_silence_gaps( + speech_segments, + non_speech_segments, + total_audio_duration=3.0, + min_silence_duration=0.5, + ) + assert len(silences) == 0 diff --git a/tests/test_embed_faces.py b/tests/test_embed_faces.py index 66f7d6e..91753a2 100644 --- a/tests/test_embed_faces.py +++ b/tests/test_embed_faces.py @@ -1,74 +1,70 @@ -import unittest -import tempfile import os import json +import pytest import numpy as np from unittest.mock import patch, MagicMock from captionalchemy.tools.cv.embed_known_faces import embed_faces -class TestEmbedFacesFunctionality(unittest.TestCase): - def setUp(self): - # Two dummy face entries - self.entries = [ - {"name": "Alice", "image_path": "path/to/alice.jpg"}, - {"name": "Bob", "image_path": "path/to/bob.jpg"}, - ] +@pytest.fixture +def face_entries(): + """Two dummy face entries.""" + return [ + {"name": "Alice", "image_path": "path/to/alice.jpg"}, + {"name": "Bob", "image_path": "path/to/bob.jpg"}, + ] - def test_embed_faces_empty_json(self): - with self.assertRaises(ValueError): - embed_faces("", "output_embeddings.json") - @patch("captionalchemy.tools.cv.embed_known_faces.FaceAnalysis") - @patch("captionalchemy.tools.cv.embed_known_faces.cv2.imread") - @patch("torch.cuda.is_available", return_value=False) - @patch("torch.backends.mps.is_available", return_value=False) - def test_embed_faces_success( - self, mock_mps, mock_cuda, mock_imread, mock_face_analysis - ): - # Prepare temporary JSON input and output paths - with tempfile.TemporaryDirectory() as tmpdir: - input_path = os.path.join(tmpdir, "known_faces.json") - output_path = os.path.join(tmpdir, "output_embeddings.json") +def test_embed_faces_empty_json(): + with pytest.raises(ValueError): + embed_faces("", "output_embeddings.json") - # Write sample known-faces JSON - for entry in self.entries: - # Create dummy image files - img_path = os.path.join(tmpdir, f"{entry['name']}.jpg") - open(img_path, "wb").close() - entry["image_path"] = img_path - with open(input_path, "w") as f: - json.dump(self.entries, f) - # Mock cv2.imread to return a dummy array for any path - mock_imread.return_value = np.random.randint( - 0, 255, (100, 100, 3), dtype=np.uint8 - ) +@patch("captionalchemy.tools.cv.embed_known_faces.FaceAnalysis") +@patch("captionalchemy.tools.cv.embed_known_faces.cv2.imread") +@patch("torch.cuda.is_available", return_value=False) +@patch("torch.backends.mps.is_available", return_value=False) +def test_embed_faces_success( + mock_mps, mock_cuda, mock_imread, mock_face_analysis, temp_dir, face_entries +): + """Test successful face embedding process.""" + input_path = os.path.join(temp_dir, "known_faces.json") + output_path = os.path.join(temp_dir, "output_embeddings.json") - # Mock FaceAnalysis to return a dummy embedding and get() method - mock_app = MagicMock() - # Each call returns a single face with a known embedding - face = MagicMock() - face.embedding = np.random.rand(512).astype(np.float32) - mock_app.get.return_value = [face] - mock_face_analysis.return_value = mock_app + # Write sample known-faces JSON + for entry in face_entries: + # Create dummy image files + img_path = os.path.join(temp_dir, f"{entry['name']}.jpg") + open(img_path, "wb").close() + entry["image_path"] = img_path + with open(input_path, "w") as f: + json.dump(face_entries, f) - # Run - embed_faces(input_path, output_path) + # Mock cv2.imread to return a dummy array for any path + mock_imread.return_value = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) - # Read output and verify - with open(output_path, "r") as f: - data = json.load(f) - # Should have two entries - self.assertEqual(len(data), 2) - names = [entry["name"] for entry in data] - self.assertCountEqual(names, ["Alice", "Bob"]) - for d in data: - emb = d["embedding"] - self.assertIsInstance(emb, list) - self.assertEqual(len(emb), 512) + # Mock FaceAnalysis to return a dummy embedding and get() method + mock_app = MagicMock() + # Each call returns a single face with a known embedding + face = MagicMock() + face.embedding = np.random.rand(512).astype(np.float32) + mock_app.get.return_value = [face] + mock_face_analysis.return_value = mock_app + # Run + embed_faces(input_path, output_path) -if __name__ == "__main__": - unittest.main() + # Read output and verify + with open(output_path, "r") as f: + data = json.load(f) + + # Should have two entries + assert len(data) == 2 + names = [entry["name"] for entry in data] + assert set(names) == {"Alice", "Bob"} + + for d in data: + emb = d["embedding"] + assert isinstance(emb, list) + assert len(emb) == 512 diff --git a/tests/test_main.py b/tests/test_main.py index c706c73..1df2a3b 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,172 +1,810 @@ -import unittest import tempfile import os import json import uuid +import pytest +import argparse from unittest.mock import Mock, patch, MagicMock, call from typing import List, Dict, Any import numpy as np import shutil +import logging +from captionalchemy.tools.audio_analysis.audio_segment_integration import ( + AudioEvent, + EventType, +) -class TestMainPipelineIntegration(unittest.TestCase): - def setUp(self): - self.test_video_url = "test_video.mp4" - self.temp_dir = tempfile.mkdtemp() - self.setup_mocks() +from captionalchemy.caption import run_pipeline, _build_arg_parser, main - def tearDown(self): - if os.path.exists(self.temp_dir): - shutil.rmtree(self.temp_dir) - def setup_mocks(self): - """Set up mocks for pipeline testing.""" - self.mock_video_manager = Mock() - self.mock_video_manager.get_video_data.return_value = ( - "/path/to/downloaded/video.mp4" +@pytest.fixture +def temp_dir(): + """Create and cleanup temporary directory.""" + temp_dir = tempfile.mkdtemp() + yield temp_dir + if os.path.exists(temp_dir): + shutil.rmtree(temp_dir) + + +@pytest.fixture +def mock_video_data(): + """Mock video-related test data.""" + return {"url": "test_video.mp4", "path": "/path/to/downloaded/video.mp4"} + + +@pytest.fixture +def mock_speech_segments(): + """Mock speech segments (VAD).""" + return [ + {"start": 0.0, "end": 5.0, "duration": 5.0}, + {"start": 10.0, "end": 15.0, "duration": 5.0}, + {"start": 20.0, "end": 25.0, "duration": 5.0}, + ] + + +@pytest.fixture +def mock_diarization_result(): + """Mock diarization result.""" + return { + "SPEAKER_00": {"start": 0.0, "end": 15.0}, + "SPEAKER_01": {"start": 10.0, "end": 25.0}, + } + + +@pytest.fixture +def mock_non_speech_events(): + """Mock non-speech events.""" + return [ + { + "start": 5.5, + "end": 8.0, + "label": "music", + "confidence": 0.9, + "duration": 2.5, + } + ] + + +@pytest.fixture +def mock_recognized_faces(): + """Mock face recognition results.""" + return [ + { + "timestamp": 0.0, + "bbox": [100, 100, 200, 200], + "face_id": "face_1", + "name": "John Doe", + }, + { + "timestamp": 10.0, + "bbox": [150, 150, 250, 250], + "face_id": "face_2", + "name": "Jane Smith", + }, + ] + + +@pytest.fixture +def mock_word_timings(): + """Mock word timings from transcriber.""" + return [ + Mock(word="Hello", start=0.0, end=0.5, duration=0.5), + Mock(word="world", start=0.6, end=1.1, duration=0.5), + Mock(word=".", start=1.1, end=1.2, duration=0.1), + Mock(word="This", start=1.3, end=1.8, duration=0.5), + Mock(word="is", start=1.9, end=2.0, duration=0.1), + Mock(word="a", start=2.1, end=2.2, duration=0.1), + Mock(word="test", start=2.3, end=2.8, duration=0.5), + ] + + +@pytest.fixture +def mock_subtitle_segments(): + """Mock subtitle segments from timing analyzer.""" + return [ + Mock( + start=0.0, + end=1.2, + text="Hello world.", + word_count=2, + char_count=12, + break_reason="sentence_ending", + ), + Mock( + start=1.3, + end=2.8, + text="This is a test.", + word_count=4, + char_count=15, + break_reason="duration_limit", + ), + ] + + +@pytest.fixture +def mock_integrated_events(mock_speech_segments): + """Mock integrated audio events.""" + events = [] + for i, segment in enumerate(mock_speech_segments): + event = Mock() + event.event_type = Mock() + event.event_type.value = "speech" + event.start = segment["start"] + event.end = segment["end"] + event.speaker_id = "SPEAKER_00" if segment["start"] < 15 else "SPEAKER_01" + events.append(event) + + # Add a non-speech event + music_event = Mock() + music_event.event_type = Mock() + music_event.event_type.value = "music" + music_event.start = 5.5 + music_event.end = 8.0 + events.append(music_event) + + return events + + +class TestRunPipeline: + """Test cases for the main run_pipeline function.""" + + @patch("caption.SRTCaptionWriter") + @patch("caption.TimingAnalyzer") + @patch("caption.Transcriber") + @patch("caption.recognize_faces") + @patch("caption.integrate_audio_segments") + @patch("caption.detect_non_speech_segments") + @patch("caption.diarize") + @patch("caption.get_speech_segments") + @patch("caption.extract_audio") + @patch("caption.VideoManager") + @patch("caption.embed_faces") + @patch("caption.whisper.load_model") + @patch("caption.torch.cuda.is_available") + def test_complete_pipeline_success( + self, + mock_cuda, + mock_whisper_load, + mock_embed_faces, + mock_video_manager_class, + mock_extract_audio, + mock_vad, + mock_diarize, + mock_non_speech, + mock_integrate_audio, + mock_recognize_faces, + mock_transcriber_class, + mock_timing_analyzer_class, + mock_srt_writer_class, + mock_video_data, + mock_speech_segments, + mock_diarization_result, + mock_non_speech_events, + mock_recognized_faces, + mock_word_timings, + mock_subtitle_segments, + mock_integrated_events, + ): + """Test the complete pipeline with successful execution.""" + + # Setup mocks + mock_cuda.return_value = True + + mock_whisper_model = Mock() + mock_whisper_load.return_value = mock_whisper_model + + mock_embed_faces.return_value = None + + mock_video_manager = Mock() + mock_video_manager.get_video_data.return_value = mock_video_data["path"] + mock_video_manager_class.return_value = mock_video_manager + + mock_extract_audio.return_value = None + + mock_vad.return_value = mock_speech_segments + mock_diarize.return_value = mock_diarization_result + mock_non_speech.return_value = mock_non_speech_events + + mock_integrate_audio.return_value = mock_integrated_events + + mock_recognize_faces.return_value = mock_recognized_faces + + mock_transcriber = Mock() + mock_transcriber.transcribe_audio.return_value = mock_word_timings + mock_transcriber_class.return_value = mock_transcriber + + mock_timing_analyzer = Mock() + mock_timing_analyzer.suggest_subtitle_segments.return_value = ( + mock_subtitle_segments ) + mock_timing_analyzer_class.return_value = mock_timing_analyzer + + # Mock SRT writer + mock_srt_writer = Mock() + mock_srt_writer_class.return_value = mock_srt_writer + + # Run the pipeline + run_pipeline( + video_url_or_path=mock_video_data["url"], + character_identification=True, + known_faces_json="test_known_faces.json", + embed_faces_json="test_embed_faces.json", + caption_output_path="test_output", + caption_format="srt", + ) + + # Verify key function calls + mock_embed_faces.assert_called_once_with( + "test_known_faces.json", "test_embed_faces.json" + ) + mock_video_manager_class.assert_called_once_with(use_file_buffer=False) + mock_whisper_load.assert_called_once() + mock_vad.assert_called_once() + mock_diarize.assert_called_once() + mock_non_speech.assert_called_once() + mock_integrate_audio.assert_called_once() + + # Verify transcriber was called for speech events + speech_events = [ + e for e in mock_integrated_events if e.event_type.value == "speech" + ] + assert mock_transcriber.transcribe_audio.call_count == len(speech_events) + + # Verify timing analyzer was called for speech events + assert mock_timing_analyzer.suggest_subtitle_segments.call_count == len( + speech_events + ) + + # Verify writer methods were called + mock_srt_writer.write.assert_called_once_with("test_output.srt") + + @patch("caption.embed_faces") + @patch("caption.VideoManager") + @patch("caption.extract_audio") + @patch("caption.get_speech_segments") + @patch("caption.torch.cuda.is_available") + @patch("caption.whisper.load_model") + def test_pipeline_no_speech_segments( + self, + mock_whisper_load, + mock_cuda, + mock_vad, + mock_extract_audio, + mock_video_manager_class, + mock_embed_faces, + mock_video_data, + ): + """Test pipeline behavior when no speech segments are detected.""" + + mock_cuda.return_value = False + mock_embed_faces.return_value = None + mock_whisper_load.return_value = Mock() + + mock_video_manager = Mock() + mock_video_manager.get_video_data.return_value = mock_video_data["path"] + mock_video_manager_class.return_value = mock_video_manager + + mock_extract_audio.return_value = None + + # Return empty speech segments + mock_vad.return_value = [] + + # This should complete without error, just log a warning + run_pipeline( + video_url_or_path=mock_video_data["url"], + character_identification=False, + ) + + # Verify early functions were called + mock_embed_faces.assert_called_once() + mock_vad.assert_called_once() + + @patch("caption.VTTCaptionWriter") + @patch("caption.embed_faces") + @patch("caption.VideoManager") + @patch("caption.extract_audio") + @patch("caption.get_speech_segments") + @patch("caption.torch.cuda.is_available") + @patch("caption.whisper.load_model") + def test_pipeline_vtt_format( + self, + mock_whisper_load, + mock_cuda, + mock_vad, + mock_extract_audio, + mock_video_manager_class, + mock_embed_faces, + mock_vtt_writer_class, + mock_video_data, + ): + """Test pipeline with VTT caption format.""" + + mock_cuda.return_value = False + mock_embed_faces.return_value = None + mock_whisper_load.return_value = Mock() + + mock_video_manager = Mock() + mock_video_manager.get_video_data.return_value = mock_video_data["path"] + mock_video_manager_class.return_value = mock_video_manager + + mock_extract_audio.return_value = None + mock_vad.return_value = [] # No speech to simplify test + + mock_vtt_writer = Mock() + mock_vtt_writer_class.return_value = mock_vtt_writer + + run_pipeline( + video_url_or_path=mock_video_data["url"], + character_identification=False, + caption_format="vtt", + ) + + # Verify VTT writer was instantiated + mock_vtt_writer_class.assert_called_once() + + @patch("caption.SAMICaptionWriter") + @patch("caption.embed_faces") + @patch("caption.VideoManager") + @patch("caption.extract_audio") + @patch("caption.get_speech_segments") + @patch("caption.torch.cuda.is_available") + @patch("caption.whisper.load_model") + def test_pipeline_sami_format( + self, + mock_whisper_load, + mock_cuda, + mock_vad, + mock_extract_audio, + mock_video_manager_class, + mock_embed_faces, + mock_sami_writer_class, + mock_video_data, + ): + """Test pipeline with SAMI caption format.""" + + mock_cuda.return_value = False + mock_embed_faces.return_value = None + mock_whisper_load.return_value = Mock() + + mock_video_manager = Mock() + mock_video_manager.get_video_data.return_value = mock_video_data["path"] + mock_video_manager_class.return_value = mock_video_manager + + mock_extract_audio.return_value = None + mock_vad.return_value = [] + + mock_sami_writer = Mock() + mock_sami_writer_class.return_value = mock_sami_writer + + run_pipeline( + video_url_or_path=mock_video_data["url"], + character_identification=False, + caption_format="smi", + ) + + mock_sami_writer_class.assert_called_once() + + def test_pipeline_with_existing_video_file(self, tmp_path): + """Test pipeline when video_url_or_path is an existing file.""" + + # Create a dummy video file + video_file = tmp_path / "test_video.mp4" + video_file.write_text("dummy video content") + + with patch("caption.embed_faces"), patch( + "caption.VideoManager" + ) as mock_vm_class, patch("caption.extract_audio"), patch( + "caption.get_speech_segments", return_value=[] + ), patch( + "caption.torch.cuda.is_available", return_value=False + ), patch( + "caption.whisper.load_model" + ): + + mock_vm = Mock() + mock_vm_class.return_value = mock_vm - self.mock_extract_audio = Mock() + run_pipeline( + video_url_or_path=str(video_file), + character_identification=False, + ) - # Mock speech segments (VAD) - self.mock_speech_segments = [ - {"start": 0.0, "end": 5.0, "duration": 5.0}, - {"start": 10.0, "end": 15.0, "duration": 5.0}, - {"start": 20.0, "end": 25.0, "duration": 5.0}, + # Should not call get_video_data since file exists + mock_vm.get_video_data.assert_not_called() + + @patch("caption.SRTCaptionWriter") + @patch("caption.TimingAnalyzer") + @patch("caption.Transcriber") + @patch("caption.integrate_audio_segments") + @patch("caption.detect_non_speech_segments") + @patch("caption.diarize") + @patch("caption.get_speech_segments") + @patch("caption.extract_audio") + @patch("caption.VideoManager") + @patch("caption.embed_faces") + @patch("caption.whisper.load_model") + @patch("caption.torch.cuda.is_available") + def test_pipeline_character_identification_disabled( + self, + mock_cuda, + mock_whisper_load, + mock_embed_faces, + mock_video_manager_class, + mock_extract_audio, + mock_vad, + mock_diarize, + mock_non_speech, + mock_integrate_audio, + mock_transcriber_class, + mock_timing_analyzer_class, + mock_srt_writer_class, + mock_speech_segments, + mock_integrated_events, + mock_word_timings, + mock_subtitle_segments, + ): + """Test pipeline with character identification disabled.""" + + # Setup basic mocks + mock_cuda.return_value = False + mock_whisper_load.return_value = Mock() + mock_embed_faces.return_value = None + + mock_vm = Mock() + mock_vm.get_video_data.return_value = "/tmp/video.mp4" + mock_video_manager_class.return_value = mock_vm + + mock_extract_audio.return_value = None + mock_vad.return_value = mock_speech_segments + mock_diarize.return_value = {"SPEAKER_00": {"start": 0.0, "end": 30.0}} + mock_non_speech.return_value = [] + + # Only speech events for this test + speech_events = [ + e for e in mock_integrated_events if e.event_type.value == "speech" ] + mock_integrate_audio.return_value = speech_events - # Mock diarization result - self.mock_diarization_result = { + mock_transcriber = Mock() + mock_transcriber.transcribe_audio.return_value = mock_word_timings + mock_transcriber_class.return_value = mock_transcriber + + mock_timing_analyzer = Mock() + mock_timing_analyzer.suggest_subtitle_segments.return_value = ( + mock_subtitle_segments + ) + mock_timing_analyzer_class.return_value = mock_timing_analyzer + + mock_writer = Mock() + mock_srt_writer_class.return_value = mock_writer + + run_pipeline( + video_url_or_path="test_video.mp4", + character_identification=False, + ) + + # embed_faces should still be called + mock_embed_faces.assert_called_once() + + # Should process speech events even without face recognition + assert mock_transcriber.transcribe_audio.call_count == len(speech_events) + + @patch("caption.SRTCaptionWriter") + @patch("caption.TimingAnalyzer") + @patch("caption.Transcriber") + @patch("caption.recognize_faces") + @patch("caption.integrate_audio_segments") + @patch("caption.detect_non_speech_segments") + @patch("caption.diarize") + @patch("caption.get_speech_segments") + @patch("caption.extract_audio") + @patch("caption.VideoManager") + @patch("caption.embed_faces") + @patch("caption.whisper.load_model") + @patch("caption.torch.cuda.is_available") + def test_pipeline_speaker_name_mapping( + self, + mock_cuda, + mock_whisper_load, + mock_embed_faces, + mock_video_manager_class, + mock_extract_audio, + mock_vad, + mock_diarize, + mock_non_speech, + mock_integrate_audio, + mock_recognize_faces, + mock_transcriber_class, + mock_timing_analyzer_class, + mock_srt_writer_class, + mock_speech_segments, + mock_word_timings, + mock_subtitle_segments, + ): + """Test that speaker names are properly mapped and reused.""" + + # Setup mocks + mock_cuda.return_value = True + mock_whisper_load.return_value = Mock() + mock_embed_faces.return_value = None + + mock_vm = Mock() + mock_vm.get_video_data.return_value = "/tmp/video.mp4" + mock_video_manager_class.return_value = mock_vm + + mock_extract_audio.return_value = None + mock_vad.return_value = mock_speech_segments + mock_diarize.return_value = { "SPEAKER_00": {"start": 0.0, "end": 15.0}, - "SPEAKER_01": {"start": 10.0, "end": 25.0}, + "SPEAKER_01": {"start": 15.0, "end": 30.0}, } + mock_non_speech.return_value = [] - # Mock non-speech events - self.mock_non_speech_events = [ - { - "start": 5.5, - "end": 8.0, - "label": "music", - "confidence": 0.9, - "duration": 2.5, - } - ] + # Create events with same speaker appearing multiple times + events = [] + for i, segment in enumerate(mock_speech_segments): + event = Mock() + event.event_type = Mock() + event.event_type.value = "speech" + event.start = segment["start"] + event.end = segment["end"] + event.speaker_id = "SPEAKER_00" # Same speaker for all + events.append(event) - # Mock face recognition results - self.mock_recognized_faces = [ - { - "timesatmp": 0.0, - "bbox": [100, 100, 200, 200], - "face_id": "face_1", - "name": "John Doe", - }, - { - "timesatmp": 10.0, - "bbox": [150, 150, 250, 250], - "face_id": "face_2", - "name": "Jane Smith", - }, - ] + mock_integrate_audio.return_value = events - # Mock word timings from transcriber - self.mock_word_timings = [ - {"word": "Hello", "start": 0.0, "end": 0.5, "duration": 0.5}, - {"word": "world", "start": 0.6, "end": 1.1, "duration": 0.5}, - {"word": ".", "start": 1.1, "end": 1.2, "duration": 0.1}, - {"word": "This", "start": 1.3, "end": 1.8, "duration": 0.5}, - {"word": "is", "start": 1.9, "end": 2.0, "duration": 0.1}, - {"word": "a", "start": 2.1, "end": 2.2, "duration": 0.1}, - {"word": "test", "start": 2.3, "end": 2.8, "duration": 0.5}, - ] + # Mock face recognition to return same name + mock_recognize_faces.return_value = [{"name": "John Doe"}] + + mock_transcriber = Mock() + mock_transcriber.transcribe_audio.return_value = mock_word_timings + mock_transcriber_class.return_value = mock_transcriber + + mock_timing_analyzer = Mock() + mock_timing_analyzer.suggest_subtitle_segments.return_value = ( + mock_subtitle_segments + ) + mock_timing_analyzer_class.return_value = mock_timing_analyzer + + mock_writer = Mock() + mock_srt_writer_class.return_value = mock_writer + + run_pipeline( + video_url_or_path="test_video.mp4", + character_identification=True, + ) + + # Face recognition should only be called once per unique speaker + assert mock_recognize_faces.call_count == 1 + + +class TestArgumentParser: + """Test cases for the command line argument parser.""" + + def test_build_arg_parser_default_values(self): + """Test that argument parser has correct default values.""" + + parser = _build_arg_parser() + args = parser.parse_args(["test_video.mp4"]) + + assert args.video == "test_video.mp4" + assert args.format == "srt" + assert args.output == "output_captions" + assert args.character_identification is True + assert args.known_faces_json == "example/known_faces.json" + assert args.embed_faces_json == "example/embed_faces.json" + assert args.verbose is False + + def test_build_arg_parser_custom_values(self): + """Test argument parser with custom values.""" + + parser = _build_arg_parser() + args = parser.parse_args( + [ + "https://example.com/video.mp4", + "--format", + "vtt", + "--output", + "my_captions", + "--no-face-id", + "--known-faces-json", + "my_faces.json", + "--embed-faces-json", + "my_embeddings.json", + "--verbose", + ] + ) + + assert args.video == "https://example.com/video.mp4" + assert args.format == "vtt" + assert args.output == "my_captions" + assert args.character_identification is False + assert args.known_faces_json == "my_faces.json" + assert args.embed_faces_json == "my_embeddings.json" + assert args.verbose is True + + def test_build_arg_parser_format_choices(self): + """Test that format argument only accepts valid choices.""" + + parser = _build_arg_parser() + + # Valid formats should work + for fmt in ["srt", "vtt", "smi"]: + args = parser.parse_args(["video.mp4", "--format", fmt]) + assert args.format == fmt - # Mock subtitle segments from timing analyzer - self.mock_subtitle_segments = [ - { - "start": 0.0, - "end": 1.2, - "text": "Hello world.", - "word_count": 2, - "char_count": 12, - "break_reason": "sentence_ending", - } + # Invalid format should raise SystemExit + with pytest.raises(SystemExit): + parser.parse_args(["video.mp4", "--format", "invalid"]) + + def test_build_arg_parser_help_text(self): + """Test that help text is properly configured.""" + + parser = _build_arg_parser() + + assert parser.prog == "captionalchemy" + assert "Download/extract audio from a video" in parser.description + + +class TestMainFunction: + """Test cases for the main() function.""" + + @patch("caption.run_pipeline") + @patch("caption.logging.getLogger") + @patch("caption.load_dotenv") + def test_main_with_default_args( + self, mock_load_dotenv, mock_get_logger, mock_run_pipeline + ): + """Test main function with default arguments.""" + + mock_logger = Mock() + mock_get_logger.return_value = mock_logger + + with patch("sys.argv", ["captionalchemy", "test_video.mp4"]): + main() + + mock_load_dotenv.assert_called_once() + mock_run_pipeline.assert_called_once_with( + video_url_or_path="test_video.mp4", + character_identification=True, + known_faces_json="example/known_faces.json", + embed_faces_json="example/embed_faces.json", + caption_output_path="output_captions", + caption_format="srt", + ) + + @patch("caption.run_pipeline") + @patch("caption.logging.getLogger") + @patch("caption.load_dotenv") + def test_main_with_verbose_logging( + self, mock_load_dotenv, mock_get_logger, mock_run_pipeline + ): + """Test main function with verbose logging enabled.""" + + mock_logger = Mock() + mock_get_logger.return_value = mock_logger + + with patch("sys.argv", ["captionalchemy", "video.mp4", "--verbose"]), patch( + "caption.logging.getLogger" + ) as mock_root_logger: + + main() + + # Should set logging level to DEBUG + mock_root_logger().setLevel.assert_called_with(logging.DEBUG) + + @patch("caption.run_pipeline") + @patch("caption.logging.getLogger") + @patch("caption.load_dotenv") + def test_main_with_custom_args( + self, mock_load_dotenv, mock_get_logger, mock_run_pipeline + ): + """Test main function with custom arguments.""" + + mock_logger = Mock() + mock_get_logger.return_value = mock_logger + + test_args = [ + "captionalchemy", + "https://example.com/video.mp4", + "--format", + "vtt", + "--output", + "custom_output", + "--no-face-id", + "--known-faces-json", + "custom_faces.json", + "--embed-faces-json", + "custom_embeddings.json", ] + with patch("sys.argv", test_args): + main() -@patch("tools.captioning.timing_analyzer.TimingAnalyzer") -@patch("tools.captioning.transcriber.Transcriber") -@patch("src.captionalchemy.tools.cv.recognize_faces.recognize_faces") -@patch("tools.audio_analysis.audio_segment_integration.integrate_audio_segments") -@patch("tools.audio_analysis.non_speech_detection.detect_non_speech_segments") -@patch("tools.audio_analysis.diarization.diarize") -@patch("tools.audio_analysis.vad.get_speech_segments") -@patch("src.captionalchemy.tools.media_utils.extract_audio.extract_audio") -@patch("src.captionalchemy.tools.media_utils.download_video.VideoManager") -@patch("src.captionalchemy.tools.cv.embed_known_faces.embed_faces") -@patch("whisper.load_model") -@patch("torch.cuda.is_available") -def test_complete_pipeline_success( - self, - mock_cuda, - mock_whisper_load, - mock_embed_faces, - mock_video_manager_class, - mock_extract_audio, - mock_vad, - mock_diarize, - mock_non_speech, - mock_integrate_audio, - mock_recognize_faces, - mock_transcriber_class, - mock_timing_analyzer_class, - mock_srt_writer_class, -): - """Test the complete pipeline with successful execution.""" - mock_cuda.return_value = True - - mock_whisper_model = Mock() - mock_whisper_load.return_value = mock_whisper_model - - mock_embed_faces.return_value = None - - mock_video_manager = Mock() - mock_video_manager.get_video_data.return_value = "/tmp/video.mp4" - mock_video_manager_class.return_value = mock_video_manager - - mock_extract_audio.return_value = None - - mock_vad.return_value = self.mock_speech_segments - mock_diarize.return_value = self.mock_diarization_result - mock_non_speech.return_value = self.mock_non_speech_events - - mock_integrated_events = [] - - for segment in self.mock_speech_segments: - mock_event = Mock() - mock_event.event_type.value = "speech" - mock_event.start = segment["start"] - mock_event.end = segment["end"] - mock_event.speaker_id = "SPEAKER_00" if segment["start"] < 15 else "SPEAKER_01" - mock_integrated_events.append(mock_event) - - mock_integrate_audio.return_value = mock_integrated_events - - # Set up mock face recognition - mock_recognize_faces.return_value = self.mock_recognized - - mock_transcriber = Mock() - mock_transcriber.transcribe_audio.return_value = [ - Mock(**wt) for wt in self.mock_word_timings - ] - mock_transcriber_class.return_value = mock_transcriber + mock_run_pipeline.assert_called_once_with( + video_url_or_path="https://example.com/video.mp4", + character_identification=False, + known_faces_json="custom_faces.json", + embed_faces_json="custom_embeddings.json", + caption_output_path="custom_output", + caption_format="vtt", + ) - mock_timing_analyzer = Mock() - mock_timing_analyzer.suggest_subtitle_segments.return_value = [ - Mock(**seg) for seg in self.mock_subtitle_segments - ] - mock_timing_analyzer_class.return_value = mock_timing_analyzer + @patch("caption.run_pipeline") + @patch("caption.logging.getLogger") + @patch("caption.load_dotenv") + def test_main_info_logging_default( + self, mock_load_dotenv, mock_get_logger, mock_run_pipeline + ): + """Test that main sets INFO logging by default.""" + + with patch("sys.argv", ["captionalchemy", "video.mp4"]), patch( + "caption.logging.getLogger" + ) as mock_root_logger: + + main() + + # Should set logging level to INFO by default + mock_root_logger().setLevel.assert_called_with(logging.INFO) + + +class TestIntegration: + """Integration tests for the complete pipeline.""" + + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}) + def test_pipeline_environment_variables(self): + """Test that environment variables are properly used.""" + + with patch("caption.embed_faces"), patch("caption.VideoManager"), patch( + "caption.extract_audio" + ), patch("caption.get_speech_segments") as mock_vad, patch( + "caption.torch.cuda.is_available", return_value=False + ), patch( + "caption.whisper.load_model" + ): + + mock_vad.return_value = [] + + run_pipeline("test_video.mp4") + + # Should pass environment token to VAD + mock_vad.assert_called_once() + call_args = mock_vad.call_args + assert call_args[0][1] == "test_token" # HF_AUTH_TOKEN + + def test_pipeline_temp_directory_cleanup(self, tmp_path): + """Test that temporary directories are properly cleaned up.""" + + # Track created temp directories + original_temp_dir = tempfile.TemporaryDirectory + temp_dirs_created = [] + + class MockTempDir: + def __init__(self): + self.name = str(tmp_path / f"temp_{len(temp_dirs_created)}") + os.makedirs(self.name, exist_ok=True) + temp_dirs_created.append(self.name) + + def __enter__(self): + return self + + def __exit__(self, *args): + if os.path.exists(self.name): + shutil.rmtree(self.name) + + with patch("tempfile.TemporaryDirectory", MockTempDir), patch( + "caption.embed_faces" + ), patch("caption.VideoManager"), patch("caption.extract_audio"), patch( + "caption.get_speech_segments", return_value=[] + ), patch( + "caption.torch.cuda.is_available", return_value=False + ), patch( + "caption.whisper.load_model" + ): - # Mock SRT writer - mock_srt_writer = Mock() - mock_srt_writer_class.return_value = mock_srt_writer + run_pipeline("test_video.mp4") - # Create simplified mock of main function + # All temp directories should be cleaned up + for temp_dir in temp_dirs_created: + assert not os.path.exists(temp_dir) diff --git a/tests/test_srt_writer.py b/tests/test_srt_writer.py index e083e70..612ff9c 100644 --- a/tests/test_srt_writer.py +++ b/tests/test_srt_writer.py @@ -1,71 +1,64 @@ -import unittest import tempfile import os +import pytest from captionalchemy.tools.captioning.writers.srt_writer import SRTCaptionWriter -class TestSRTWriter(unittest.TestCase): - def test_add_caption_text_and_silence_music(self): - writer = SRTCaptionWriter() +def test_add_caption_text_and_silence_music(): + writer = SRTCaptionWriter() - # Speech with text - writer.add_caption( - 0.0, 1.0, speaker="SPEAKER_00", text="Hello world.", event_type="speech" - ) - # Short silence (should skip) - writer.add_caption( - 1.0, 1.2, speaker="SPEAKER_00", text="", event_type="silence" - ) - # Long silence (should add) - writer.add_caption( - 1.2, 3.0, speaker="SPEAKER_00", text="", event_type="silence" - ) - # Music event (should skip) - writer.add_caption(3.0, 4.5, event_type="music") - # Long music event (should add) - writer.add_caption(4.5, 6.0, event_type="music") + # Speech with text + writer.add_caption( + 0.0, 1.0, speaker="SPEAKER_00", text="Hello world.", event_type="speech" + ) + # Short silence (should skip) + writer.add_caption(1.0, 1.2, speaker="SPEAKER_00", text="", event_type="silence") + # Long silence (should add) + writer.add_caption(1.2, 3.0, speaker="SPEAKER_00", text="", event_type="silence") + # Music event (should skip) + writer.add_caption(3.0, 4.5, event_type="music") + # Long music event (should add) + writer.add_caption(4.5, 6.0, event_type="music") - caps = writer._captions - # Expect 1 speech, 1 silence, and 1 music caption, 3 entries total - self.assertEqual(len(caps), 3) - self.assertEqual(caps[0].text, "Hello world.") - self.assertEqual(caps[1]["text"], "[SILENCE]") - self.assertEqual(caps[2]["text"], "[MUSIC PLAYING]") + caps = writer._captions + # Expect 1 speech, 1 silence, and 1 music caption, 3 entries total + assert len(caps) == 3 + assert caps[0].text == "Hello world." + assert caps[1]["text"] == "[SILENCE]" + assert caps[2]["text"] == "[MUSIC PLAYING]" - def test_write_srt_file(self): - writer = SRTCaptionWriter() - # Two entires with same speaker - writer.add_caption( - 0.0, 1.0, speaker="SPEAKER_00", text="Hello world.", event_type="speech" - ) - writer.add_caption( - 1.0, 2.0, speaker="SPEAKER_00", text="This is a test.", event_type="speech" - ) - # Different speaker - writer.add_caption( - 2.0, 3.0, speaker="SPEAKER_01", text="Goodbye.", event_type="speech" - ) +def test_write_srt_file(): + writer = SRTCaptionWriter() - tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".srt") - tmp.close() - path = tmp.name + # Two entries with same speaker + writer.add_caption( + 0.0, 1.0, speaker="SPEAKER_00", text="Hello world.", event_type="speech" + ) + writer.add_caption( + 1.0, 2.0, speaker="SPEAKER_00", text="This is a test.", event_type="speech" + ) + # Different speaker + writer.add_caption( + 2.0, 3.0, speaker="SPEAKER_01", text="Goodbye.", event_type="speech" + ) - try: - writer.write(path) - content = open(path, "r", encoding="utf-8").read().splitlines() - # Check index and timings in first block - self.assertEqual(content[0], "1") - self.assertTrue("-->" in content[1]) - # Speaker line for first caption - self.assertTrue("SPEAKER_00: Hello world." in content[2]) - # Next blocks exist - self.assertIn("2", content) - self.assertIn("3", content) - finally: - os.remove(path) + with tempfile.NamedTemporaryFile(delete=False, suffix=".srt") as tmp: + path = tmp.name + try: + writer.write(path) + with open(path, "r", encoding="utf-8") as f: + content = f.read().splitlines() -if __name__ == "__main__": - unittest.main() + # Check index and timings in first block + assert content[0] == "1" + assert "-->" in content[1] + # Speaker line for first caption + assert "SPEAKER_00: Hello world." in content[2] + # Next blocks exist + assert "2" in content + assert "3" in content + finally: + os.remove(path) diff --git a/tests/test_transcriber.py b/tests/test_transcriber.py index 2c67a0d..5352ca4 100644 --- a/tests/test_transcriber.py +++ b/tests/test_transcriber.py @@ -1,7 +1,7 @@ -import unittest import tempfile import os import subprocess +import pytest from unittest.mock import patch, Mock import numpy as np import torch @@ -9,106 +9,110 @@ from captionalchemy.tools.captioning.transcriber import Transcriber, WordTiming -class TestWordTimingDataclass(unittest.TestCase): - def test_word_timing_creation(self): - """Test creating a WordTiming instance.""" - word_timing = WordTiming(word="test", start=0.0, end=1.0, is_subword=True) - self.assertEqual(word_timing.word, "test") - self.assertEqual(word_timing.start, 0.0) - self.assertEqual(word_timing.end, 1.0) - self.assertFalse(word_timing.is_punctuation) - self.assertFalse(word_timing.is_sentence_ending) - self.assertTrue(word_timing.is_subword) - - def test_word_timing_invalid(self): - with self.assertRaises(ValueError): - WordTiming(word="", start=0.0, end=1.0) # empty word - with self.assertRaises(ValueError): - WordTiming(word="hi", start=-1.0, end=0.5) # negative start - with self.assertRaises(ValueError): - WordTiming(word="hi", start=1.0, end=0.5) # end < start - - -class TestTranscriberParse(unittest.TestCase): - def setUp(self): - self.transcriber = Transcriber() - - def test_parse_timestamps(self): - """Test parsing timestamps from a Whisper cpp output.""" - self.assertAlmostEqual( - self.transcriber._parse_timestamps("00:01:23.456"), 83.456, places=3 - ) - self.assertAlmostEqual( - self.transcriber._parse_timestamps("2:03:04.007"), - 2 * 3600 + 3 * 60 + 4.007, - places=3, - ) - with self.assertRaises(ValueError): - self.transcriber._parse_timestamps("99:99:99") - - def test_parse_line(self): - """Test parsing line from Whisper cpp output.""" - line = "[00:00:00.000 --> 00:00:00.500] Hello" - wt = self.transcriber._parse_line(line) - self.assertEqual(wt.word, "Hello") - self.assertAlmostEqual(wt.start, 0.0, places=3) - self.assertAlmostEqual(wt.end, 0.5, places=3) - # punctuation - line2 = "[00:00:01.000 --> 00:00:01.200] ." - wt2 = self.transcriber._parse_line(line2) - self.assertEqual(wt2.word, ".") - self.assertTrue(wt2.is_punctuation) - - @patch("tempfile.NamedTemporaryFile") - @patch("subprocess.run") - @patch("whisper.load_model") - @patch("os.remove") - def test_transcribe_audio_python_api( - self, mock_remove, mock_load, mock_run, mock_tempfile - ): - """Test transcribing audio using Whisper Python API.""" - # Mock temp file - fake = Mock() - fake.name = "/tmp/fake.wav" - mock_tempfile.return_value = fake - - # Mock ffmpeg trim - mock_run.return_value = Mock(returncode=0) - - # Mock whisper model load & transcribe - mock_model = Mock() - mock_load.return_value = mock_model - mock_model.transcribe.return_value = { - "text": "Hello world.", - "segments": [ - { - "words": [ - {"word": "Hello", "start": 0.0, "end": 0.5}, - {"word": "world", "start": 0.6, "end": 1.1}, - ] - } - ], - } - - transcriber = Transcriber() - words = transcriber.transcribe_audio( - audio_file="in.wav", - start=1.0, - end=2.0, - model="base", - whisper_build_path=None, - whisper_model_path=None, - platform="linux", - ) - - self.assertEqual(len(words), 2) - self.assertEqual(words[0].word, "Hello") - self.assertAlmostEqual(words[0].start, 0.0, places=3) - self.assertAlmostEqual(words[0].end, 0.5, places=3) - self.assertEqual(words[1].word, "world") - self.assertAlmostEqual(words[1].start, 0.6, places=3) - self.assertAlmostEqual(words[1].end, 1.1, places=3) - - -if __name__ == "__main__": - unittest.main() +def test_word_timing_creation(): + """Test creating a WordTiming instance.""" + word_timing = WordTiming(word="test", start=0.0, end=1.0, is_subword=True) + assert word_timing.word == "test" + assert word_timing.start == 0.0 + assert word_timing.end == 1.0 + assert word_timing.is_punctuation is False + assert word_timing.is_sentence_ending is False + assert word_timing.is_subword is True + + +def test_word_timing_invalid_empty_word(): + with pytest.raises(ValueError): + WordTiming(word="", start=0.0, end=1.0) # empty word + + +def test_word_timing_invalid_negative_start(): + with pytest.raises(ValueError): + WordTiming(word="hi", start=-1.0, end=0.5) # negative start + + +def test_word_timing_invalid_end_before_start(): + with pytest.raises(ValueError): + WordTiming(word="hi", start=1.0, end=0.5) # end < start + + +@pytest.fixture +def transcriber(): + return Transcriber() + + +def test_parse_timestamps(transcriber): + """Test parsing timestamps from a Whisper cpp output.""" + result = transcriber._parse_timestamps("00:01:23.456") + assert abs(result - 83.456) < 1e-3 + + result = transcriber._parse_timestamps("2:03:04.007") + expected = 2 * 3600 + 3 * 60 + 4.007 + assert abs(result - expected) < 1e-3 + + with pytest.raises(ValueError): + transcriber._parse_timestamps("99:99:99") + + +def test_parse_line(transcriber): + """Test parsing line from Whisper cpp output.""" + line = "[00:00:00.000 --> 00:00:00.500] Hello" + wt = transcriber._parse_line(line) + assert wt.word == "Hello" + assert abs(wt.start - 0.0) < 1e-3 + assert abs(wt.end - 0.5) < 1e-3 + + # punctuation + line2 = "[00:00:01.000 --> 00:00:01.200] ." + wt2 = transcriber._parse_line(line2) + assert wt2.word == "." + assert wt2.is_punctuation is True + + +@patch("tempfile.NamedTemporaryFile") +@patch("subprocess.run") +@patch("whisper.load_model") +@patch("os.remove") +def test_transcribe_audio_python_api( + mock_remove, mock_load, mock_run, mock_tempfile, transcriber +): + """Test transcribing audio using Whisper Python API.""" + # Mock temp file + fake = Mock() + fake.name = "/tmp/fake.wav" + mock_tempfile.return_value = fake + + # Mock ffmpeg trim + mock_run.return_value = Mock(returncode=0) + + # Mock whisper model load & transcribe + mock_model = Mock() + mock_load.return_value = mock_model + mock_model.transcribe.return_value = { + "text": "Hello world.", + "segments": [ + { + "words": [ + {"word": "Hello", "start": 0.0, "end": 0.5}, + {"word": "world", "start": 0.6, "end": 1.1}, + ] + } + ], + } + + words = transcriber.transcribe_audio( + audio_file="in.wav", + start=1.0, + end=2.0, + model="base", + whisper_build_path=None, + whisper_model_path=None, + platform="linux", + ) + + assert len(words) == 2 + assert words[0].word == "Hello" + assert abs(words[0].start - 0.0) < 1e-3 + assert abs(words[0].end - 0.5) < 1e-3 + assert words[1].word == "world" + assert abs(words[1].start - 0.6) < 1e-3 + assert abs(words[1].end - 1.1) < 1e-3 From 570303871a56e6a634461beea1b34507bac763da Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Mon, 9 Jun 2025 21:54:58 -0400 Subject: [PATCH 04/21] adding python version --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bae5806..ea99f6e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - python-version: [3.8, 3.9, 3.10, 3.11] + python-version: [3.8, 3.9, 3.10.0, 3.11, 3.12] os: [ubuntu-latest, macos-latest, windows-latest] steps: - name: Checkout code From 694ea700fa7035511a852b1bf3a3f2be15c0d708 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Sat, 14 Jun 2025 20:45:15 -0400 Subject: [PATCH 05/21] changing python version --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ea99f6e..c3eeb3a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - python-version: [3.8, 3.9, 3.10.0, 3.11, 3.12] + python-version: [3.8, 3.9, 3.10, 3.11, 3.12] os: [ubuntu-latest, macos-latest, windows-latest] steps: - name: Checkout code From 439d5f237833deec3b04405d157daa34a5355d36 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Tue, 17 Jun 2025 13:59:51 -0400 Subject: [PATCH 06/21] requirements.dev --- .github/workflows/ci.yml | 2 +- requirements-dev.txt | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) create mode 100644 requirements-dev.txt diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c3eeb3a..42ab9cf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - python-version: [3.8, 3.9, 3.10, 3.11, 3.12] + python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] os: [ubuntu-latest, macos-latest, windows-latest] steps: - name: Checkout code diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..c378bac --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,3 @@ +pytest==8.4.0 +mypy==1.16.0 +flake8==7.2.0 From ffb217319442c58afb04069a87f7f37eaf27426a Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Tue, 17 Jun 2025 14:12:22 -0400 Subject: [PATCH 07/21] fix install typo --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 42ab9cf..888065a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -49,7 +49,7 @@ jobs: sudo apt-get install -y ffmpeg libsm6 libxext6 libfontconfig1 libxrender1 # Install audio libraries for pydub sudo apt-get install -y libavcodec-extra - pip intall torch==${{ env.PYTORCH_VERSION }} torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu + pip install torch==${{ env.PYTORCH_VERSION }} torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu - name: Install dependencies (macOS) if: runner.os == 'macOS' From 977c0b0f73f12103b77b0c2dfb25135338210e78 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Tue, 17 Jun 2025 14:24:05 -0400 Subject: [PATCH 08/21] removing specified pytorch version --- .github/workflows/ci.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 888065a..8e1b93d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,7 +15,7 @@ on: workflow_dispatch: env: - PYTORCH_VERSION: "2.7.0" + PYTORCH_VERSION: "2.7" jobs: build_and_test: @@ -49,14 +49,14 @@ jobs: sudo apt-get install -y ffmpeg libsm6 libxext6 libfontconfig1 libxrender1 # Install audio libraries for pydub sudo apt-get install -y libavcodec-extra - pip install torch==${{ env.PYTORCH_VERSION }} torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu + pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu - name: Install dependencies (macOS) if: runner.os == 'macOS' run: | brew update brew install ffmpeg - pip3 install torch==${{ env.PYTORCH_VERSION }} torchvision torchaudio + pip3 install torch torchvision torchaudio - name: Install dependencies (Windows) if: runner.os == 'Windows' @@ -65,7 +65,7 @@ jobs: choco install -y ffmpeg # Add ffmpeg to PATH echo "C:\ProgramData\chocolatey\lib\ffmpeg\tools\ffmpeg\bin" >> $GITHUB_PATH - pip3 install torch==${{ env.PYTORCH_VERSION }} torchvision torchaudio + pip3 install torch torchvision torchaudio - name: Upgrade pip and dependencies run: | From 5acfdcc8fa83c9c4c9c6ea158cb56e02886148dc Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Tue, 17 Jun 2025 21:04:24 -0400 Subject: [PATCH 09/21] removing testing for python 3.8 and 3.9 --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8e1b93d..9f8ed75 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + python-version: ["3.10", "3.11", "3.12"] os: [ubuntu-latest, macos-latest, windows-latest] steps: - name: Checkout code From ad4d177e0cdc139e78051cadc485f86b1e79c0e8 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Wed, 18 Jun 2025 11:59:36 -0400 Subject: [PATCH 10/21] removing unneeded imports from test files --- tests/test_main.py | 14 +++++--------- tests/test_srt_writer.py | 1 - tests/test_transcriber.py | 4 ---- 3 files changed, 5 insertions(+), 14 deletions(-) diff --git a/tests/test_main.py b/tests/test_main.py index 1df2a3b..a5053f2 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,19 +1,15 @@ import tempfile import os -import json -import uuid import pytest -import argparse -from unittest.mock import Mock, patch, MagicMock, call -from typing import List, Dict, Any +from unittest.mock import Mock, patch import numpy as np import shutil import logging -from captionalchemy.tools.audio_analysis.audio_segment_integration import ( - AudioEvent, - EventType, -) +# from captionalchemy.tools.audio_analysis.audio_segment_integration import ( +# AudioEvent, +# EventType, +# ) from captionalchemy.caption import run_pipeline, _build_arg_parser, main diff --git a/tests/test_srt_writer.py b/tests/test_srt_writer.py index 612ff9c..184519a 100644 --- a/tests/test_srt_writer.py +++ b/tests/test_srt_writer.py @@ -1,6 +1,5 @@ import tempfile import os -import pytest from captionalchemy.tools.captioning.writers.srt_writer import SRTCaptionWriter diff --git a/tests/test_transcriber.py b/tests/test_transcriber.py index 5352ca4..d8de129 100644 --- a/tests/test_transcriber.py +++ b/tests/test_transcriber.py @@ -1,10 +1,6 @@ import tempfile -import os -import subprocess import pytest from unittest.mock import patch, Mock -import numpy as np -import torch from captionalchemy.tools.captioning.transcriber import Transcriber, WordTiming From ce985f3b17119244a0703c8c227263d3c972e794 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Wed, 18 Jun 2025 12:01:51 -0400 Subject: [PATCH 11/21] one more --- tests/test_transcriber.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_transcriber.py b/tests/test_transcriber.py index d8de129..eff4aaf 100644 --- a/tests/test_transcriber.py +++ b/tests/test_transcriber.py @@ -1,4 +1,3 @@ -import tempfile import pytest from unittest.mock import patch, Mock From 098d42c6a58ae84f18cddb3385aff7143f04b493 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Tue, 1 Jul 2025 21:01:59 -0400 Subject: [PATCH 12/21] fixing linter --- tests/test_main.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_main.py b/tests/test_main.py index a5053f2..3a0642a 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -2,7 +2,6 @@ import os import pytest from unittest.mock import Mock, patch -import numpy as np import shutil import logging @@ -773,7 +772,6 @@ def test_pipeline_temp_directory_cleanup(self, tmp_path): """Test that temporary directories are properly cleaned up.""" # Track created temp directories - original_temp_dir = tempfile.TemporaryDirectory temp_dirs_created = [] class MockTempDir: From 593811748da614a751f2b8b72130ab15ae5a2760 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Tue, 1 Jul 2025 21:18:31 -0400 Subject: [PATCH 13/21] commenting out type checking for now --- .github/workflows/ci.yml | 9 +- .gitignore | 1 + CLAUDE.md | 79 +++++++++++++++ README.md | 207 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 292 insertions(+), 4 deletions(-) create mode 100644 CLAUDE.md diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9f8ed75..3e9addc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -77,10 +77,11 @@ jobs: run: | flake8 src/ tests/ - - name: Run type checks with mypy - run: | - mypy --install-types --non-interactive --ignore-missing-imports - mypy src/ + # Will implement type checking later + # - name: Run type checks with mypy + # run: | + # mypy --install-types --non-interactive --ignore-missing-imports + # mypy src/ - name: Run tests run: | pytest --cov=captionalchemy -v diff --git a/.gitignore b/.gitignore index 117a1da..2790271 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ venv .venv __pycache__ +.mypy_cache *.env diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..b51fb7d --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,79 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +CaptionAlchemy is a Python package for creating closed captions with face detection and recognition. It combines audio transcription, speaker diarization, and facial recognition to generate accurate subtitles with speaker identification. + +## Architecture + +The package follows a modular structure under `src/captionalchemy/`: + +- **Main Pipeline**: `caption.py` - Core orchestration of the caption generation process +- **Audio Analysis**: `tools/audio_analysis/` - Voice activity detection, speaker diarization, and non-speech detection +- **Captioning**: `tools/captioning/` - Whisper transcription, timing analysis, and output writers (SRT, VTT, SAMI) +- **Computer Vision**: `tools/cv/` - Face embedding and recognition for speaker identification +- **Media Utils**: `tools/media_utils/` - Video download and audio extraction utilities + +## Key Components + +### Core Pipeline Flow + +1. Embed known faces from JSON configuration +2. Download/extract audio from video +3. Run Voice Activity Detection (VAD) and speaker diarization +4. Transcribe audio segments with Whisper +5. Identify speakers using facial recognition +6. Generate timestamped captions in chosen format + +### Dependencies + +- **Audio Processing**: pyannote.audio, pydub, librosa, openai-whisper +- **Computer Vision**: opencv-python, insightface, onnxruntime +- **Deep Learning**: Uses CUDA when available, falls back to CPU + +## Development Commands + +### Installation + +```bash +pip install -e . +``` + +### Testing + +```bash +pytest +``` + +### Linting + +```bash +flake8 +mypy src/ +``` + +### Running the Package + +```bash +captionalchemy -f srt -o output_captions +``` + +## Configuration + +- **Environment**: Uses `.env` file for configuration (HF_AUTH_TOKEN for Hugging Face models) +- **Face Recognition**: Requires `known_faces.json`. `embed_faces.json` is an artifact generated. +- **Whisper**: Integrates with whisper.cpp for transcription (requires models in `whisper.cpp/models/`) + +## External Dependencies + +- Whisper.cpp integration for transcription +- Hugging Face models for audio analysis +- ONNX runtime for face recognition models + +## Code Style + +- Line length: 110 characters (flake8 configured) +- Python 3.10+ compatible +- Type hints encouraged diff --git a/README.md b/README.md index e69de29..9449711 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,207 @@ +# CaptionAlchemy + +A Python package for creating intelligent closed captions with face detection and speaker recognition. + +## Features + +- **Audio Transcription**: Powered by OpenAI Whisper for high-quality speech-to-text +- **Speaker Diarization**: Identifies different speakers in audio +- **Face Recognition**: Links speakers to known faces for character identification +- **Multiple Output Formats**: Supports SRT, VTT, and SAMI caption formats +- **Voice Activity Detection**: Intelligently detects speech vs non-speech segments +- **GPU Acceleration**: Automatic CUDA support when available + +## Installation + +```bash +pip install -e . +``` + +### Prerequisites + +- Python 3.10+ +- FFmpeg (for video/audio processing) +- CUDA-capable GPU (optional, for acceleration) +- Whisper.cpp capable (optional on MacOS) + +## Quick Start + +1. **Set up environment variables** (create `.env` file): + + ``` + HF_AUTH_TOKEN=your_huggingface_token_here + ``` + +2. **Prepare known faces** (optional, for speaker identification): + Create `known_faces.json`: + + ```json + [ + { + "name": "Speaker Name", + "image_path": "path/to/speaker/photo.jpg" + } + ] + ``` + +3. **Generate captions**: + ```bash + captionalchemy video.mp4 -f srt -o my_captions + ``` + +## Usage + +### Basic Usage + +```bash +# Generate SRT captions from video file +captionalchemy video.mp4 + +# Generate VTT captions from YouTube URL +captionalchemy "https://youtube.com/watch?v=VIDEO_ID" -f vtt -o output + +# Disable face recognition +captionalchemy video.mp4 --no-face-id +``` + +### Command Line Options + +``` +captionalchemy VIDEO [OPTIONS] + +Arguments: + VIDEO Video file path or URL + +Options: + -f, --format Caption format: srt, vtt, smi (default: srt) + -o, --output Output file base name (default: output_captions) + --no-face-id Disable face recognition + --known-faces-json Path to known faces JSON (default: example/known_faces.json) + --embed-faces-json Path to face embeddings JSON (default: example/embed_faces.json) + -v, --verbose Enable debug logging +``` + +## How It Works + +1. **Face Embedding**: Pre-processes known faces into embeddings +2. **Audio Extraction**: Extracts audio from video files +3. **Voice Activity Detection**: Identifies speech segments +4. **Speaker Diarization**: Separates different speakers +5. **Transcription**: Converts speech to text using Whisper +6. **Face Recognition**: Matches speakers to known faces (if enabled) +7. **Caption Generation**: Creates timestamped captions with speaker names + +## Configuration + +### Known Faces Setup + +Create a `known_faces.json` file with speaker information: + +```json +[ + { + "name": "John Doe", + "image_path": "photos/john_doe.jpg" + }, + { + "name": "Jane Smith", + "image_path": "photos/jane_smith.png" + } +] +``` + +### Environment Variables + +- `HF_AUTH_TOKEN`: Hugging Face token for accessing pyannote models + +## Output Examples + +### SRT Format + +``` +1 +00:00:03,254 --> 00:00:06,890 +John Doe: Welcome to our presentation on quantum computing. + +2 +00:00:07,120 --> 00:00:10,456 +Jane Smith: Thanks John. Let's start with the basics. +``` + +### VTT Format + +``` +WEBVTT + +00:03.254 --> 00:06.890 +John Doe: Welcome to our presentation on quantum computing. + +00:07.120 --> 00:10.456 +Jane Smith: Thanks John. Let's start with the basics. +``` + +## Development + +### Setup Development Environment + +```bash +# Install in development mode +pip install -e . + +# Install development dependencies +pip install -r requirements-dev.txt +``` + +### Running Tests + +```bash +pytest +``` + +### Code Quality + +```bash +# Linting +flake8 + +# Type checking +mypy src/ +``` + +## Requirements + +See `requirements.txt` for the complete list of dependencies. Key packages include: + +- `openai-whisper`: Speech transcription +- `pyannote.audio`: Speaker diarization +- `opencv-python`: Computer vision +- `insightface`: Face recognition +- `torch`: Deep learning framework + +## License + +MIT License - see LICENSE file for details. + +## Contributing + +1. Fork the repository +2. Create a feature branch +3. Make your changes +4. Add tests for new functionality +5. Run the test suite +6. Submit a pull request + +## Troubleshooting + +### Common Issues + +- **CUDA out of memory**: Use CPU-only mode or reduce batch sizes +- **Missing models**: Ensure whisper.cpp models are downloaded +- **Face recognition errors**: Verify image paths in known_faces.json +- **Audio extraction fails**: Check that FFmpeg is installed + +### Getting Help + +- Check the logs with `-v` flag for detailed error information +- Ensure all dependencies are properly installed +- Verify video file format compatibility From abc979bdff3ac6e57a7db7e027a1053bef557916 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Tue, 1 Jul 2025 21:54:11 -0400 Subject: [PATCH 14/21] updates to gh actions, pyproject.toml --- .github/workflows/ci.yml | 18 +++++++++++------- pyproject.toml | 16 +++++++++++++++- requirements-dev.txt | 1 + src/captionalchemy/caption.py | 14 ++++++++------ 4 files changed, 35 insertions(+), 14 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3e9addc..b6f9889 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -77,6 +77,10 @@ jobs: run: | flake8 src/ tests/ + - name: Format code with black + run: | + black --check src/ tests/ + # Will implement type checking later # - name: Run type checks with mypy # run: | @@ -90,6 +94,11 @@ jobs: needs: build_and_test if: startsWith(github.ref, 'refs/tags/v') runs-on: ubuntu-latest + permissions: + id-token: write + environment: + name: pypi + url: https://pypi.org/p/captionalchemy steps: - name: Checkout code uses: actions/checkout@v4 @@ -99,15 +108,10 @@ jobs: with: python-version: "3.11" - - name: Upgrade pip and build + - name: Build distributions run: | python -m pip install --upgrade pip build python -m build - name: Publish to PyPI - env: - TWINE_USERNAME: __token__ - TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} - run: | - python -m pip install --upgrade twine - twine upload --skip-existing dist/* + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/pyproject.toml b/pyproject.toml index 4d9f532..8f69fb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,13 @@ description = "A Python package to create closed captions with face detection an readme = "README.md" license = { text = "MIT" } authors = [{ name = "Ben Batman", email = "benbatman2@gmail.com" }] +keywords = [ + "caption", + "closed captions", + "face detection", + "face recognition", + "video processing", +] dependencies = [ "pyannote.audio==3.3.2", @@ -25,13 +32,20 @@ dependencies = [ "tqdm==4.67.1", ] +[project.optional-dependencies] +dev = ["pytest==8.4.0", "mypy==1.16.0", "flake8==7.2.0", "black==25.1.0"] + classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", ] -requires-python = ">=3.8" +[tool.setuptools.packages.find] +where = ["src"] +include = ["captionalchemy*"] + +requires-python = [">=3.10"] [project.scripts] captionalchemy = "captionalchemy.caption:main" diff --git a/requirements-dev.txt b/requirements-dev.txt index c378bac..6a5b87a 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,3 +1,4 @@ pytest==8.4.0 mypy==1.16.0 flake8==7.2.0 +black==25.1.0 diff --git a/src/captionalchemy/caption.py b/src/captionalchemy/caption.py index 6eae8fe..fa62039 100644 --- a/src/captionalchemy/caption.py +++ b/src/captionalchemy/caption.py @@ -10,19 +10,21 @@ import whisper import torch -from tools.audio_analysis.diarization import diarize +from captionalchemy.tools.audio_analysis.diarization import diarize from captionalchemy.tools.media_utils.download_video import VideoManager from captionalchemy.tools.cv.recognize_faces import recognize_faces from captionalchemy.tools.media_utils.extract_audio import extract_audio -from tools.captioning.transcriber import Transcriber -from tools.captioning.timing_analyzer import TimingAnalyzer +from captionalchemy.tools.captioning.transcriber import Transcriber +from captionalchemy.tools.captioning.timing_analyzer import TimingAnalyzer from captionalchemy.tools.cv.embed_known_faces import embed_faces from captionalchemy.tools.captioning.writers.srt_writer import SRTCaptionWriter from captionalchemy.tools.captioning.writers.vtt_writer import VTTCaptionWriter from captionalchemy.tools.captioning.writers.sami_writer import SAMICaptionWriter -from tools.audio_analysis.vad import get_speech_segments -from tools.audio_analysis.non_speech_detection import detect_non_speech_segments -from tools.audio_analysis.audio_segment_integration import ( +from captionalchemy.tools.audio_analysis.vad import get_speech_segments +from captionalchemy.tools.audio_analysis.non_speech_detection import ( + detect_non_speech_segments, +) +from captionalchemy.tools.audio_analysis.audio_segment_integration import ( integrate_audio_segments, ) From c02b1d6c3cd789832a36a52154cf7f080c17d2f3 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Wed, 2 Jul 2025 08:21:12 -0400 Subject: [PATCH 15/21] update .gitignore --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 2790271..363aa80 100644 --- a/.gitignore +++ b/.gitignore @@ -4,8 +4,9 @@ __pycache__ .mypy_cache *.env - +.env # Whisper whisper.cpp /whisper.cpp + From cc123c82bee474bc609211d8d453189d8f39cc40 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Wed, 2 Jul 2025 09:07:50 -0400 Subject: [PATCH 16/21] small update to github actions --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b6f9889..0cb5ecd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -98,7 +98,7 @@ jobs: id-token: write environment: name: pypi - url: https://pypi.org/p/captionalchemy + url: https://pypi.org/project/captionalchemy steps: - name: Checkout code uses: actions/checkout@v4 From a6e3bea06502464dc0c11cabc52b6b3dbf7a81d6 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Thu, 3 Jul 2025 10:06:17 -0400 Subject: [PATCH 17/21] new unit tests and integration tests --- .coverage | Bin 0 -> 53248 bytes .github/workflows/ci.yml | 5 +- .gitignore | 1 + pyproject.toml | 14 +- requirements-dev.txt | 1 + src/captionalchemy/caption.py | 13 +- .../audio_segment_integration.py | 5 +- .../tools/captioning/writers/sami_writer.py | 3 +- tests/test_audio_integration.py | 2 +- tests/test_diarization.py | 352 ++++++++++++++ tests/test_main.py | 197 ++++---- tests/test_sami_writer.py | 330 +++++++++++++ tests/test_srt_writer.py | 8 +- tests/test_transcriber.py | 20 +- tests/test_vtt_writer.py | 442 ++++++++++++++++++ 15 files changed, 1277 insertions(+), 116 deletions(-) create mode 100644 .coverage create mode 100644 tests/test_diarization.py create mode 100644 tests/test_sami_writer.py create mode 100644 tests/test_vtt_writer.py diff --git a/.coverage b/.coverage new file mode 100644 index 0000000000000000000000000000000000000000..e575f9271e600d8959c2a0ae0e881b1272dba5ac GIT binary patch literal 53248 zcmeI4eQX>@6~K4z*0;C!HFlG5O+)cYsCC@fKCDWG5ZWY#sH$i&NgOH?CELBbvAxaR z?qzq+iBp>R5~K*Es)9d2mC{gEMHN&6{((XTq^JdnA~`J;5{Qq45492bXp7RMA;#go z{qWs6c2vZgD&dWMGrO}h^XC2L&CJcMXEtxV&v3aqV%b&QRc{wM1yK~%sj46d3VgEg z37rfVQg{QQ*mAwlr6LTN&u7`^giPX5f$hi+vHr~A?5*inGgalw=_9hPl;8n5B0vO) z01^2ACeYZJk@L%!iw7QY^BoCLHkc__bJRlzUUX2luDs{E{W&A@1p< z+a>r*w`l_rO;zq4fr3hwUB=oMZQU%5ai>q!XEiKYP|rQCudQuW3$CSB4AWE7bPU(9 zOqK89rMk<@^D5Yf&_W0BTD5_)9z=y6^m7KXTIe)0&yzDq9%LSY%9&P(wV4q`LuRke zX1!Vjo*G*R=v=-*el{)0Rzr+nO^(jxFlvs+-2HI-hO)KDC*% z&r`#+`czZ`EZ~-OBENcxD1cS+JyQm>?#iGP_o9kt&j47;Hf_>yx)cdrZcD{ z@*BDrGM!ly&<1BOQ;p>bIlp3sctG*YC}^!9klRcsY{_|s(eP1~>rUMc2a>A|l3|YK zC>i#V;0dVHFDLTrS1hDXkWL#6E|LvdlJm=!i4DI;;JAdJjy5GQrr8DNOa$+}`Ze7h zQ?IN$MnlDWIuiLc%ND9QRzVv?KXGtd%&YRD(2_4NyA-iPcVRj-#Ey7y--$1z}9#c9j= zco?hpjOh-JLQaobw5$r(O@D5H9QQgP({NO;4p+`fYnx>_6IjORc#l`7aB{(6ei%Ak z*uUVSTGh>RFkj&`?3W5-KXy)(;}0T`QSE_F=|C5qQ)XCb@DxX^u>64Qz0gc5oYOGm<$EoAxpU11@~qX|)q%*Mb9YeF-?Y`7K!V zDTJwl?>eboT;PlBk^q0?hyW2F0z`la5CI}U1c(3;AOb{y2oQnmk3d|Ei88MLW9-iY zdlNizM1Tko0U|&IhyW2F0z`la5CI}U1c<=3Ngxx8-^{{)@%RW6<16yPKLLDnVD-RY z3bcx`lL9-*-nljvq*@aJB0vO)01+SpM1Tko0U|&IhyW2F0wIA+{AMxy6(E+1uTX*? z0r35QSMptfeVeV#y_oxa_Eh$f?6S-YnN8`_>HXWS1U<>$(0APG4lKm>>Y5g-CY zfCyaE1ctjLq0wL1;=mr$La|uXUD%jf*j%!Wn(GwG#;DP6n5NZV(rdU)yx-ddYS{(H zE){|}y@I$^!t%^iiV2)7ehzx~pjHFP%8v`f~ z2c49)ohB&IM3smb)tuo1Ou;tx!a+Eei4D^cv9y-{3tQ)*8~X#*UT%Z5F$tt?iISl| z660tAW^Fn!RKw9RAZc$LBuzvpX%0pI*5RX7Zn~PcbKVY_8ZnUlz~af)K%U`xbIY&@ zqN19yW=qBQ|6R$80(&-jG4nIFF85M)pE9KUB6~WwHPx49sUz7Bq|anFrzfxJ>7FW2 z1c(3;AOb{y2oM25*djIh#Yn!mwXFZU#-)Z8O{xFM`oEJ)jr$i-tiS$eyQBt>QcBDE zVgB`ht|m1~i!8Bi{h$3G2y9#bXDU)-Ta+4JZT+8)+DztM|EH=_WBX#MbjA8#8I~HW zqLJ2I|0hSK#>Qw0z3Td3jvCdRm23O@KM}RGmi~L?`oF`F8e1Y|xX}7vdRS^qL@8+w z#jCCV;}1!V2NqFw%lbbCmfo*KG-J)K|LOcc`Z7cWhyW2F0z`la5CI}U1c(3;AOb|- zx+V}8<3bkS|BGx|fIo6XfCvx)B0vO)01+SpM1Tko0U|&Ih`@D7ARf;u`2N2$cTiv# z+1uM2UtfFd!r`Yc{OQ+g^7kA+ z(6=-TDQ?S-oc!)u;S0b2`(u;CJwrF`Uz)+#j32x8#AB1^PAh+9avFhlzCm`a+#MxV?-kw=D)X^a-jSqFq{Nt!F ztqC)yFV8$D2#MMYNkK(vN%X*-fTgz}tP@1pi@MW`1TP045pQ1UMql{YfCQ!QkY={e zT=<4KnOG{Uo4oUxB@Z`VUVY1@q06E$b$Nf@%kb7)|JpiT5^st_`a9xNZ@hb9weXW2 ze>nWa*Sm#Ro_god^jEqQF?g1Xi5;S(K#1@EJK0wSc8Q&1XV~lPIQuPomA%YfWItx# zXWwDZuqWB$kboQ!AOb{y2oM1xKm>>Y5g-CYfCvx)A`l?~{3Su`?DPW0ydanJf^60c zG8r#Od%rn|sg!r8C|;0EdVwr^K_cM=9UWdENnQ|-dqFHF!4Cl5`+xiuf*cVb0z`la z5CI}U1c(3;AOb{y2oM1xaQza%_y4s1zkb`7G$R5;fCvx)B0vO)01+SpM1Tko0V03| zaQ}Z7dlvuy|2&)kaGJfr{=)tUdjMWzQ|y=QC3ciO5BmVV2T8~g0U|&IhyW2F0z`la z5CI}U1c(3;AOed|Al?P%^~XEWW9a43%c7S-FO6ObJq5iadNO(m^g7U!(2JuNLl4jY F|2Il>C0hUh literal 0 HcmV?d00001 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0cb5ecd..a3a0b1d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -69,9 +69,8 @@ jobs: - name: Upgrade pip and dependencies run: | - python -m pip install --upgrade pip setuptools wheel - pip install -r requirements.txt - pip install -r requirements-dev.txt + python -m pip install --upgrade pip setuptools wheel + pip install -e .[dev] - name: Lint with flake8 run: | diff --git a/.gitignore b/.gitignore index 363aa80..9607835 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ venv .venv __pycache__ .mypy_cache +captionalchemy.egg-info *.env .env diff --git a/pyproject.toml b/pyproject.toml index 8f69fb7..6bd3817 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,20 +32,26 @@ dependencies = [ "tqdm==4.67.1", ] -[project.optional-dependencies] -dev = ["pytest==8.4.0", "mypy==1.16.0", "flake8==7.2.0", "black==25.1.0"] - classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", ] +requires-python = ">=3.10" + + +[project.optional-dependencies] +dev = ["pytest==8.4.0", "mypy==1.16.0", "flake8==7.2.0", "black==25.1.0"] + [tool.setuptools.packages.find] where = ["src"] include = ["captionalchemy*"] -requires-python = [">=3.10"] [project.scripts] captionalchemy = "captionalchemy.caption:main" + + +[tool.pytest.ini_options] +pythonpath = ["src"] diff --git a/requirements-dev.txt b/requirements-dev.txt index 6a5b87a..a6e1bea 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -2,3 +2,4 @@ pytest==8.4.0 mypy==1.16.0 flake8==7.2.0 black==25.1.0 +pytest-cov==6.2.1 diff --git a/src/captionalchemy/caption.py b/src/captionalchemy/caption.py index fa62039..29df031 100644 --- a/src/captionalchemy/caption.py +++ b/src/captionalchemy/caption.py @@ -46,6 +46,14 @@ def run_pipeline( 3. Runs VAD + diarization, 4. Runs Whisper transcription + face-based speaker ID, 5. Writes captions in the chosen format (SRT, VTT, SMI). + + Args: + video_url_or_path: URL or local path of the video to caption. + character_identification: Enable/disable face-based speaker identification. + known_faces_json: Path to JSON file listing known faces to embed. + embed_faces_json: JSON path to store face embeddings. + caption_output_path: Base path (without extension) for output captions. + caption_format: Format for output captions (srt, vtt, smi). """ logger.info("Embedding known faces...") embed_faces(known_faces_json, embed_faces_json) @@ -98,9 +106,6 @@ def run_pipeline( diarization_result = diarize( audio_path ) # { "SPEAKER_00": {"start": 3.25409375, "end": 606.2990937500001}, ..., } - # diarization_result = { - # "SPEAKER_00": {"start": 3.25409375, "end": 606.2990937500001} - # } logger.info("Completed diarization.") logger.debug(f"Diarization result: {diarization_result}") @@ -151,7 +156,7 @@ def run_pipeline( # Map speaker ID to name speaker_id_to_name[speaker_id] = speaker_name else: - speaker_name = speaker_id_to_name[speaker_id] + speaker_name = speaker_id_to_name.get(speaker_id, speaker_id) # Add speaker name to the audio event audio_event.speaker_name = speaker_name diff --git a/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py b/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py index eca6308..259c6e6 100644 --- a/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py +++ b/src/captionalchemy/tools/audio_analysis/audio_segment_integration.py @@ -64,6 +64,9 @@ def assign_speakers_to_speech_segment( # Find which speaker's time range this segment falls into # Looking for the speaker whose range has the maximum overlap with this segment for speaker_id, time_range in diarization.items(): + print( + f"Checking speaker {speaker_id} for segment {segment_start} - {segment_end}" + ) speaker_start = time_range["start"] speaker_end = time_range["end"] @@ -72,7 +75,7 @@ def assign_speakers_to_speech_segment( overlap_end = min(segment_end, speaker_end) overlap = max(0, overlap_end - overlap_start) - # If this egment is mostly within this speaker's range, assign it + # If this segment is mostly within this speaker's range, assign it segment_duration = segment_end - segment_start overlap_ratio = overlap / segment_duration if segment_duration > 0 else 0 diff --git a/src/captionalchemy/tools/captioning/writers/sami_writer.py b/src/captionalchemy/tools/captioning/writers/sami_writer.py index 7b3c6c0..cfbde82 100644 --- a/src/captionalchemy/tools/captioning/writers/sami_writer.py +++ b/src/captionalchemy/tools/captioning/writers/sami_writer.py @@ -203,7 +203,8 @@ def _escape_text(self, text: str) -> str: Escape HTML entities and handle special characters for SAMI format. """ # First escape basic HTML entities - escaped = html.escape(text) + # Single quotes get outputted to the numeric ', not the HTML decimal ' + escaped = html.escape(text, quote=False) # Replace other common characters that might cause issues escaped = escaped.replace('"', """) diff --git a/tests/test_audio_integration.py b/tests/test_audio_integration.py index 95ce350..57f5e83 100644 --- a/tests/test_audio_integration.py +++ b/tests/test_audio_integration.py @@ -131,7 +131,7 @@ def test_no_silence_when_too_short(): silences = identify_silence_gaps( speech_segments, non_speech_segments, - total_audio_duration=3.0, + total_audio_duration=2.1, # Ensure the final silence check is valid min_silence_duration=0.5, ) assert len(silences) == 0 diff --git a/tests/test_diarization.py b/tests/test_diarization.py new file mode 100644 index 0000000..8c02c29 --- /dev/null +++ b/tests/test_diarization.py @@ -0,0 +1,352 @@ +import pytest +import os +import tempfile +from unittest.mock import patch, MagicMock +from types import SimpleNamespace + +from captionalchemy.tools.audio_analysis.diarization import diarize + + +class TestDiarization: + """Test suite for diarization functionality.""" + + @patch.dict(os.environ, {}, clear=True) + def test_diarize_missing_hf_token(self): + """Test that diarize raises ValueError when HF_AUTH_TOKEN is missing.""" + with pytest.raises( + ValueError, match="Hugging Face authentication token not found" + ): + diarize("dummy_audio.wav") + + @patch.dict(os.environ, {"HF_AUTH_TOKEN": ""}, clear=True) + def test_diarize_empty_hf_token(self): + """Test that diarize raises ValueError when HF_AUTH_TOKEN is empty.""" + with pytest.raises( + ValueError, match="Hugging Face authentication token not found" + ): + diarize("dummy_audio.wav") + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_success_cpu(self, mock_torch, mock_pipeline_class): + """Test successful diarization on CPU.""" + # Mock torch.cuda.is_available() to return False (CPU) + mock_torch.cuda.is_available.return_value = False + mock_torch.device.return_value = "cpu" + + # Mock pipeline instance and its methods + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + # Mock diarization result + mock_turn1 = SimpleNamespace(start=1.0, end=5.0) + mock_turn2 = SimpleNamespace(start=6.0, end=10.0) + mock_turn3 = SimpleNamespace(start=11.0, end=15.0) + + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [ + (mock_turn1, None, "SPEAKER_00"), + (mock_turn2, None, "SPEAKER_01"), + (mock_turn3, None, "SPEAKER_00"), # Same speaker continues + ] + mock_pipeline.return_value = mock_diarization + + result = diarize("test_audio.wav") + + # Verify pipeline was created correctly + mock_pipeline_class.from_pretrained.assert_called_once_with( + "pyannote/speaker-diarization-3.1", use_auth_token="test_token" + ) + + # Verify pipeline was moved to device + mock_pipeline.to.assert_called_once_with("cpu") + + # Verify pipeline was called with audio path + mock_pipeline.assert_called_once_with("test_audio.wav") + + # Verify result structure + expected_result = { + "SPEAKER_00": { + "start": 1.0, + "end": 15.0, + }, # First speaker spans 1-5 and 11-15 + "SPEAKER_01": {"start": 6.0, "end": 10.0}, # Second speaker spans 6-10 + } + assert result == expected_result + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_success_gpu(self, mock_torch, mock_pipeline_class): + """Test successful diarization on GPU.""" + # Mock torch.cuda.is_available() to return True (GPU) + mock_torch.cuda.is_available.return_value = True + mock_device = MagicMock() + mock_torch.device.return_value = mock_device + + # Mock pipeline instance + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + # Mock diarization result + mock_turn = SimpleNamespace(start=2.5, end=8.7) + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [ + (mock_turn, None, "SPEAKER_00"), + ] + mock_pipeline.return_value = mock_diarization + + result = diarize("test_audio.wav") + + # Verify GPU device was used + mock_torch.device.assert_called_once_with("cuda") + mock_pipeline.to.assert_called_once_with(mock_device) + + # Verify result + expected_result = {"SPEAKER_00": {"start": 2.5, "end": 8.7}} + assert result == expected_result + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_multiple_speakers_complex(self, mock_torch, mock_pipeline_class): + """Test diarization with multiple speakers and complex timing.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.device.return_value = "cpu" + + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + # Create complex scenario with overlapping and multiple segments + mock_turns = [ + SimpleNamespace(start=0.0, end=3.0), # SPEAKER_00 + SimpleNamespace(start=2.5, end=5.0), # SPEAKER_01 (overlaps) + SimpleNamespace(start=5.5, end=8.0), # SPEAKER_00 (continues) + SimpleNamespace(start=8.5, end=12.0), # SPEAKER_02 (new speaker) + SimpleNamespace(start=12.5, end=15.0), # SPEAKER_01 (returns) + SimpleNamespace(start=15.5, end=18.0), # SPEAKER_00 (returns) + ] + + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [ + (mock_turns[0], None, "SPEAKER_00"), + (mock_turns[1], None, "SPEAKER_01"), + (mock_turns[2], None, "SPEAKER_00"), + (mock_turns[3], None, "SPEAKER_02"), + (mock_turns[4], None, "SPEAKER_01"), + (mock_turns[5], None, "SPEAKER_00"), + ] + mock_pipeline.return_value = mock_diarization + + result = diarize("complex_audio.wav") + + # The function extends end times for repeated speakers + expected_result = { + "SPEAKER_00": { + "start": 0.0, + "end": 18.0, + }, # Start at 0.0, extended to final segment + "SPEAKER_01": { + "start": 2.5, + "end": 15.0, + }, # Start at 2.5, extended to final segment + "SPEAKER_02": {"start": 8.5, "end": 12.0}, # Only one segment + } + assert result == expected_result + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_single_speaker(self, mock_torch, mock_pipeline_class): + """Test diarization with a single speaker.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.device.return_value = "cpu" + + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + # Single speaker with multiple segments + mock_turns = [ + SimpleNamespace(start=1.0, end=5.0), + SimpleNamespace(start=7.0, end=12.0), + SimpleNamespace(start=15.0, end=20.0), + ] + + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [ + (mock_turns[0], None, "SPEAKER_00"), + (mock_turns[1], None, "SPEAKER_00"), + (mock_turns[2], None, "SPEAKER_00"), + ] + mock_pipeline.return_value = mock_diarization + + result = diarize("single_speaker_audio.wav") + + expected_result = {"SPEAKER_00": {"start": 1.0, "end": 20.0}} + assert result == expected_result + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_no_speech(self, mock_torch, mock_pipeline_class): + """Test diarization with no speech detected.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.device.return_value = "cpu" + + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + # No speech segments + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [] + mock_pipeline.return_value = mock_diarization + + result = diarize("silent_audio.wav") + + assert result == {} + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_pipeline_exception(self, mock_torch, mock_pipeline_class): + """Test handling of pipeline exceptions.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.device.return_value = "cpu" + + # Mock pipeline to raise an exception + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + mock_pipeline.side_effect = RuntimeError("Pipeline processing failed") + + with pytest.raises(RuntimeError, match="Pipeline processing failed"): + diarize("problematic_audio.wav") + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_pipeline_creation_exception(self, mock_torch, mock_pipeline_class): + """Test handling of pipeline creation exceptions.""" + mock_torch.cuda.is_available.return_value = False + + # Mock pipeline creation to fail + mock_pipeline_class.from_pretrained.side_effect = Exception( + "Failed to load model" + ) + + with pytest.raises(Exception, match="Failed to load model"): + diarize("audio.wav") + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_with_real_file_path(self, mock_torch, mock_pipeline_class): + """Test diarization with a real file path.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.device.return_value = "cpu" + + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + mock_turn = SimpleNamespace(start=0.0, end=10.0) + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [ + (mock_turn, None, "SPEAKER_00"), + ] + mock_pipeline.return_value = mock_diarization + + # Use a real temporary file path + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file: + temp_path = temp_file.name + + try: + result = diarize(temp_path) + + # Verify the actual file path was passed to the pipeline + mock_pipeline.assert_called_once_with(temp_path) + + expected_result = {"SPEAKER_00": {"start": 0.0, "end": 10.0}} + assert result == expected_result + + finally: + os.unlink(temp_path) + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "hf_test123"}, clear=True) + def test_diarize_auth_token_passed_correctly(self, mock_torch, mock_pipeline_class): + """Test that the auth token is passed correctly to the pipeline.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.device.return_value = "cpu" + + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [] + mock_pipeline.return_value = mock_diarization + + diarize("test.wav") + + # Verify the correct auth token was used + mock_pipeline_class.from_pretrained.assert_called_once_with( + "pyannote/speaker-diarization-3.1", use_auth_token="hf_test123" + ) + + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_fractional_timestamps(self, mock_torch, mock_pipeline_class): + """Test diarization with fractional timestamps.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.device.return_value = "cpu" + + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + # Use fractional timestamps + mock_turns = [ + SimpleNamespace(start=1.234, end=5.678), + SimpleNamespace(start=6.789, end=10.123), + ] + + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [ + (mock_turns[0], None, "SPEAKER_00"), + (mock_turns[1], None, "SPEAKER_01"), + ] + mock_pipeline.return_value = mock_diarization + + result = diarize("fractional_audio.wav") + + expected_result = { + "SPEAKER_00": {"start": 1.234, "end": 5.678}, + "SPEAKER_01": {"start": 6.789, "end": 10.123}, + } + assert result == expected_result + + @patch("captionalchemy.tools.audio_analysis.diarization.logger") + @patch("captionalchemy.tools.audio_analysis.diarization.Pipeline") + @patch("captionalchemy.tools.audio_analysis.diarization.torch") + @patch.dict(os.environ, {"HF_AUTH_TOKEN": "test_token"}, clear=True) + def test_diarize_logging(self, mock_torch, mock_pipeline_class, mock_logger): + """Test that appropriate logging messages are generated.""" + mock_torch.cuda.is_available.return_value = True + mock_device = MagicMock() + mock_device.__str__ = MagicMock(return_value="cuda:0") + mock_torch.device.return_value = mock_device + + mock_pipeline = MagicMock() + mock_pipeline_class.from_pretrained.return_value = mock_pipeline + + mock_diarization = MagicMock() + mock_diarization.itertracks.return_value = [] + mock_pipeline.return_value = mock_diarization + + diarize("test_audio.wav") + + # Check that logging messages were called + mock_logger.info.assert_any_call( + "Diarizing audio file: test_audio.wav on cuda:0..." + ) + mock_logger.info.assert_any_call("Diarization completed.") diff --git a/tests/test_main.py b/tests/test_main.py index 3a0642a..cc83766 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -144,19 +144,19 @@ def mock_integrated_events(mock_speech_segments): class TestRunPipeline: """Test cases for the main run_pipeline function.""" - @patch("caption.SRTCaptionWriter") - @patch("caption.TimingAnalyzer") - @patch("caption.Transcriber") - @patch("caption.recognize_faces") - @patch("caption.integrate_audio_segments") - @patch("caption.detect_non_speech_segments") - @patch("caption.diarize") - @patch("caption.get_speech_segments") - @patch("caption.extract_audio") - @patch("caption.VideoManager") - @patch("caption.embed_faces") - @patch("caption.whisper.load_model") - @patch("caption.torch.cuda.is_available") + @patch("captionalchemy.caption.SRTCaptionWriter") + @patch("captionalchemy.caption.TimingAnalyzer") + @patch("captionalchemy.caption.Transcriber") + @patch("captionalchemy.caption.recognize_faces") + @patch("captionalchemy.caption.integrate_audio_segments") + @patch("captionalchemy.caption.detect_non_speech_segments") + @patch("captionalchemy.caption.diarize") + @patch("captionalchemy.caption.get_speech_segments") + @patch("captionalchemy.caption.extract_audio") + @patch("captionalchemy.caption.VideoManager") + @patch("captionalchemy.caption.embed_faces") + @patch("captionalchemy.caption.whisper.load_model") + @patch("captionalchemy.caption.torch.cuda.is_available") def test_complete_pipeline_success( self, mock_cuda, @@ -254,12 +254,13 @@ def test_complete_pipeline_success( # Verify writer methods were called mock_srt_writer.write.assert_called_once_with("test_output.srt") - @patch("caption.embed_faces") - @patch("caption.VideoManager") - @patch("caption.extract_audio") - @patch("caption.get_speech_segments") - @patch("caption.torch.cuda.is_available") - @patch("caption.whisper.load_model") + @patch("captionalchemy.caption.detect_non_speech_segments") + @patch("captionalchemy.caption.embed_faces") + @patch("captionalchemy.caption.VideoManager") + @patch("captionalchemy.caption.extract_audio") + @patch("captionalchemy.caption.get_speech_segments") + @patch("captionalchemy.caption.torch.cuda.is_available") + @patch("captionalchemy.caption.whisper.load_model") def test_pipeline_no_speech_segments( self, mock_whisper_load, @@ -268,6 +269,7 @@ def test_pipeline_no_speech_segments( mock_extract_audio, mock_video_manager_class, mock_embed_faces, + mock_detect_non_speech, mock_video_data, ): """Test pipeline behavior when no speech segments are detected.""" @@ -284,6 +286,7 @@ def test_pipeline_no_speech_segments( # Return empty speech segments mock_vad.return_value = [] + mock_detect_non_speech.return_value = [] # This should complete without error, just log a warning run_pipeline( @@ -295,13 +298,14 @@ def test_pipeline_no_speech_segments( mock_embed_faces.assert_called_once() mock_vad.assert_called_once() - @patch("caption.VTTCaptionWriter") - @patch("caption.embed_faces") - @patch("caption.VideoManager") - @patch("caption.extract_audio") - @patch("caption.get_speech_segments") - @patch("caption.torch.cuda.is_available") - @patch("caption.whisper.load_model") + @patch("captionalchemy.caption.detect_non_speech_segments") + @patch("captionalchemy.caption.VTTCaptionWriter") + @patch("captionalchemy.caption.embed_faces") + @patch("captionalchemy.caption.VideoManager") + @patch("captionalchemy.caption.extract_audio") + @patch("captionalchemy.caption.get_speech_segments") + @patch("captionalchemy.caption.torch.cuda.is_available") + @patch("captionalchemy.caption.whisper.load_model") def test_pipeline_vtt_format( self, mock_whisper_load, @@ -311,6 +315,7 @@ def test_pipeline_vtt_format( mock_video_manager_class, mock_embed_faces, mock_vtt_writer_class, + mock_detect_non_speech, mock_video_data, ): """Test pipeline with VTT caption format.""" @@ -325,6 +330,7 @@ def test_pipeline_vtt_format( mock_extract_audio.return_value = None mock_vad.return_value = [] # No speech to simplify test + mock_detect_non_speech.return_value = [] mock_vtt_writer = Mock() mock_vtt_writer_class.return_value = mock_vtt_writer @@ -338,13 +344,14 @@ def test_pipeline_vtt_format( # Verify VTT writer was instantiated mock_vtt_writer_class.assert_called_once() - @patch("caption.SAMICaptionWriter") - @patch("caption.embed_faces") - @patch("caption.VideoManager") - @patch("caption.extract_audio") - @patch("caption.get_speech_segments") - @patch("caption.torch.cuda.is_available") - @patch("caption.whisper.load_model") + @patch("captionalchemy.caption.detect_non_speech_segments") + @patch("captionalchemy.caption.SAMICaptionWriter") + @patch("captionalchemy.caption.embed_faces") + @patch("captionalchemy.caption.VideoManager") + @patch("captionalchemy.caption.extract_audio") + @patch("captionalchemy.caption.get_speech_segments") + @patch("captionalchemy.caption.torch.cuda.is_available") + @patch("captionalchemy.caption.whisper.load_model") def test_pipeline_sami_format( self, mock_whisper_load, @@ -354,6 +361,7 @@ def test_pipeline_sami_format( mock_video_manager_class, mock_embed_faces, mock_sami_writer_class, + mock_detect_non_speech, mock_video_data, ): """Test pipeline with SAMI caption format.""" @@ -368,6 +376,7 @@ def test_pipeline_sami_format( mock_extract_audio.return_value = None mock_vad.return_value = [] + mock_detect_non_speech.return_value = [] mock_sami_writer = Mock() mock_sami_writer_class.return_value = mock_sami_writer @@ -387,14 +396,16 @@ def test_pipeline_with_existing_video_file(self, tmp_path): video_file = tmp_path / "test_video.mp4" video_file.write_text("dummy video content") - with patch("caption.embed_faces"), patch( - "caption.VideoManager" - ) as mock_vm_class, patch("caption.extract_audio"), patch( - "caption.get_speech_segments", return_value=[] + with patch("captionalchemy.caption.embed_faces"), patch( + "captionalchemy.caption.VideoManager" + ) as mock_vm_class, patch("captionalchemy.caption.extract_audio"), patch( + "captionalchemy.caption.get_speech_segments", return_value=[] ), patch( - "caption.torch.cuda.is_available", return_value=False + "captionalchemy.caption.detect_non_speech_segments", return_value=[] ), patch( - "caption.whisper.load_model" + "captionalchemy.caption.torch.cuda.is_available", return_value=False + ), patch( + "captionalchemy.caption.whisper.load_model" ): mock_vm = Mock() @@ -408,18 +419,18 @@ def test_pipeline_with_existing_video_file(self, tmp_path): # Should not call get_video_data since file exists mock_vm.get_video_data.assert_not_called() - @patch("caption.SRTCaptionWriter") - @patch("caption.TimingAnalyzer") - @patch("caption.Transcriber") - @patch("caption.integrate_audio_segments") - @patch("caption.detect_non_speech_segments") - @patch("caption.diarize") - @patch("caption.get_speech_segments") - @patch("caption.extract_audio") - @patch("caption.VideoManager") - @patch("caption.embed_faces") - @patch("caption.whisper.load_model") - @patch("caption.torch.cuda.is_available") + @patch("captionalchemy.caption.SRTCaptionWriter") + @patch("captionalchemy.caption.TimingAnalyzer") + @patch("captionalchemy.caption.Transcriber") + @patch("captionalchemy.caption.integrate_audio_segments") + @patch("captionalchemy.caption.detect_non_speech_segments") + @patch("captionalchemy.caption.diarize") + @patch("captionalchemy.caption.get_speech_segments") + @patch("captionalchemy.caption.extract_audio") + @patch("captionalchemy.caption.VideoManager") + @patch("captionalchemy.caption.embed_faces") + @patch("captionalchemy.caption.whisper.load_model") + @patch("captionalchemy.caption.torch.cuda.is_available") def test_pipeline_character_identification_disabled( self, mock_cuda, @@ -485,19 +496,19 @@ def test_pipeline_character_identification_disabled( # Should process speech events even without face recognition assert mock_transcriber.transcribe_audio.call_count == len(speech_events) - @patch("caption.SRTCaptionWriter") - @patch("caption.TimingAnalyzer") - @patch("caption.Transcriber") - @patch("caption.recognize_faces") - @patch("caption.integrate_audio_segments") - @patch("caption.detect_non_speech_segments") - @patch("caption.diarize") - @patch("caption.get_speech_segments") - @patch("caption.extract_audio") - @patch("caption.VideoManager") - @patch("caption.embed_faces") - @patch("caption.whisper.load_model") - @patch("caption.torch.cuda.is_available") + @patch("captionalchemy.caption.SRTCaptionWriter") + @patch("captionalchemy.caption.TimingAnalyzer") + @patch("captionalchemy.caption.Transcriber") + @patch("captionalchemy.caption.recognize_faces") + @patch("captionalchemy.caption.integrate_audio_segments") + @patch("captionalchemy.caption.detect_non_speech_segments") + @patch("captionalchemy.caption.diarize") + @patch("captionalchemy.caption.get_speech_segments") + @patch("captionalchemy.caption.extract_audio") + @patch("captionalchemy.caption.VideoManager") + @patch("captionalchemy.caption.embed_faces") + @patch("captionalchemy.caption.whisper.load_model") + @patch("captionalchemy.caption.torch.cuda.is_available") def test_pipeline_speaker_name_mapping( self, mock_cuda, @@ -645,9 +656,9 @@ def test_build_arg_parser_help_text(self): class TestMainFunction: """Test cases for the main() function.""" - @patch("caption.run_pipeline") - @patch("caption.logging.getLogger") - @patch("caption.load_dotenv") + @patch("captionalchemy.caption.run_pipeline") + @patch("captionalchemy.caption.logging.getLogger") + @patch("captionalchemy.caption.load_dotenv") def test_main_with_default_args( self, mock_load_dotenv, mock_get_logger, mock_run_pipeline ): @@ -669,9 +680,9 @@ def test_main_with_default_args( caption_format="srt", ) - @patch("caption.run_pipeline") - @patch("caption.logging.getLogger") - @patch("caption.load_dotenv") + @patch("captionalchemy.caption.run_pipeline") + @patch("captionalchemy.caption.logging.getLogger") + @patch("captionalchemy.caption.load_dotenv") def test_main_with_verbose_logging( self, mock_load_dotenv, mock_get_logger, mock_run_pipeline ): @@ -681,7 +692,7 @@ def test_main_with_verbose_logging( mock_get_logger.return_value = mock_logger with patch("sys.argv", ["captionalchemy", "video.mp4", "--verbose"]), patch( - "caption.logging.getLogger" + "captionalchemy.caption.logging.getLogger" ) as mock_root_logger: main() @@ -689,9 +700,9 @@ def test_main_with_verbose_logging( # Should set logging level to DEBUG mock_root_logger().setLevel.assert_called_with(logging.DEBUG) - @patch("caption.run_pipeline") - @patch("caption.logging.getLogger") - @patch("caption.load_dotenv") + @patch("captionalchemy.caption.run_pipeline") + @patch("captionalchemy.caption.logging.getLogger") + @patch("captionalchemy.caption.load_dotenv") def test_main_with_custom_args( self, mock_load_dotenv, mock_get_logger, mock_run_pipeline ): @@ -726,16 +737,16 @@ def test_main_with_custom_args( caption_format="vtt", ) - @patch("caption.run_pipeline") - @patch("caption.logging.getLogger") - @patch("caption.load_dotenv") + @patch("captionalchemy.caption.run_pipeline") + @patch("captionalchemy.caption.logging.getLogger") + @patch("captionalchemy.caption.load_dotenv") def test_main_info_logging_default( self, mock_load_dotenv, mock_get_logger, mock_run_pipeline ): """Test that main sets INFO logging by default.""" with patch("sys.argv", ["captionalchemy", "video.mp4"]), patch( - "caption.logging.getLogger" + "captionalchemy.caption.logging.getLogger" ) as mock_root_logger: main() @@ -751,12 +762,16 @@ class TestIntegration: def test_pipeline_environment_variables(self): """Test that environment variables are properly used.""" - with patch("caption.embed_faces"), patch("caption.VideoManager"), patch( - "caption.extract_audio" - ), patch("caption.get_speech_segments") as mock_vad, patch( - "caption.torch.cuda.is_available", return_value=False + with patch("captionalchemy.caption.embed_faces"), patch( + "captionalchemy.caption.VideoManager" + ), patch("captionalchemy.caption.extract_audio"), patch( + "captionalchemy.caption.get_speech_segments" + ) as mock_vad, patch( + "captionalchemy.caption.detect_non_speech_segments", return_value=[] + ), patch( + "captionalchemy.caption.torch.cuda.is_available", return_value=False ), patch( - "caption.whisper.load_model" + "captionalchemy.caption.whisper.load_model" ): mock_vad.return_value = [] @@ -781,20 +796,24 @@ def __init__(self): temp_dirs_created.append(self.name) def __enter__(self): - return self + return self.name def __exit__(self, *args): if os.path.exists(self.name): shutil.rmtree(self.name) with patch("tempfile.TemporaryDirectory", MockTempDir), patch( - "caption.embed_faces" - ), patch("caption.VideoManager"), patch("caption.extract_audio"), patch( - "caption.get_speech_segments", return_value=[] + "captionalchemy.caption.embed_faces" + ), patch("captionalchemy.caption.VideoManager"), patch( + "captionalchemy.caption.extract_audio" + ), patch( + "captionalchemy.caption.get_speech_segments", return_value=[] + ), patch( + "captionalchemy.caption.detect_non_speech_segments", return_value=[] ), patch( - "caption.torch.cuda.is_available", return_value=False + "captionalchemy.caption.torch.cuda.is_available", return_value=False ), patch( - "caption.whisper.load_model" + "captionalchemy.caption.whisper.load_model" ): run_pipeline("test_video.mp4") diff --git a/tests/test_sami_writer.py b/tests/test_sami_writer.py new file mode 100644 index 0000000..0d8fe27 --- /dev/null +++ b/tests/test_sami_writer.py @@ -0,0 +1,330 @@ +import tempfile +import os +from unittest.mock import patch, mock_open + +from captionalchemy.tools.captioning.writers.sami_writer import SAMICaptionWriter + + +class TestSAMICaptionWriter: + """Test suite for SAMICaptionWriter class.""" + + def test_init_default_values(self): + """Test initialization with default values.""" + writer = SAMICaptionWriter() + + assert writer.min_silence_duration == 1.0 + assert writer.min_music_duration == 1.5 + assert writer.language == "en-US" + assert writer.title == "Generated Captions" + assert writer.last_speaker is None + assert writer._captions == [] + + def test_init_custom_values(self): + """Test initialization with custom values.""" + writer = SAMICaptionWriter( + min_silence_duration=2.0, + min_music_duration=3.0, + language="fr-FR", + title="Custom Title", + ) + + assert writer.min_silence_duration == 2.0 + assert writer.min_music_duration == 3.0 + assert writer.language == "fr-FR" + assert writer.title == "Custom Title" + + def test_add_caption_with_text(self): + """Test adding a caption with text.""" + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, + speaker="John Doe", + text="Hello world", + event_type="speech", + ) + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["start"] == 1.0 + assert caption["end"] == 3.0 + assert caption["speaker"] == "John Doe" + assert caption["text"] == "Hello world" + assert caption["event_type"] == "speech" + + def test_add_caption_with_newlines(self): + """Test adding a caption with newlines that should be replaced.""" + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, end=3.0, text="Hello\nworld\ntest", event_type="speech" + ) + + caption = writer._captions[0] + assert caption["text"] == "Hello world test" + + def test_add_caption_music_long_duration(self): + """Test adding a music caption with long duration.""" + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, # 2 seconds, > min_music_duration (1.5) + event_type="music", + ) + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["text"] == "[music playing]" + + def test_add_caption_music_short_duration(self): + """Test adding a music caption with short duration (should be skipped).""" + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, + end=2.0, # 1 second, < min_music_duration (1.5) + event_type="music", + ) + + assert len(writer._captions) == 0 + + def test_add_caption_silence_long_duration(self): + """Test adding a silence caption with long duration.""" + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, # 2 seconds, > min_silence_duration (1.0) + event_type="silence", + ) + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["text"] == "[silence]" + + def test_add_caption_silence_short_duration(self): + """Test adding a silence caption with short duration (should be skipped).""" + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, + end=1.5, # 0.5 seconds, < min_silence_duration (1.0) + event_type="silence", + ) + + assert len(writer._captions) == 0 + + def test_add_caption_with_label(self): + """Test adding a caption with a label.""" + writer = SAMICaptionWriter() + writer.add_caption(start=1.0, end=3.0, event_type="other", label="APPLAUSE") + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["text"] == "[applause]" + + def test_add_caption_fallback(self): + """Test adding a caption with fallback text.""" + writer = SAMICaptionWriter() + writer.add_caption(start=1.0, end=3.0, event_type="unknown") + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["text"] == "[audio]" + + def test_escape_text(self): + """Test the _escape_text method.""" + writer = SAMICaptionWriter() + + # Test HTML entities + assert writer._escape_text("Hello & world") == "Hello & world" + assert writer._escape_text("Hello < world") == "Hello < world" + assert writer._escape_text("Hello > world") == "Hello > world" + + # Test quotes + assert writer._escape_text('Hello "world"') == "Hello "world"" + assert writer._escape_text("Hello 'world'") == "Hello 'world'" + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_empty_captions(self, mock_makedirs, mock_file): + """Test writing with no captions.""" + writer = SAMICaptionWriter() + writer.write("test.smi") + + mock_makedirs.assert_called_once_with(".", exist_ok=True) + mock_file.assert_called_once_with("test.smi", "w", encoding="utf-8") + + # Check that the basic SAMI structure is written + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + assert "" in written_content + assert "" in written_content + assert "" in written_content + assert "" in written_content + assert "" in written_content + assert "" in written_content + assert "Generated Captions" in written_content + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_with_captions(self, mock_makedirs, mock_file): + """Test writing with captions.""" + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, + speaker="John Doe", + text="Hello world", + event_type="speech", + ) + writer.add_caption(start=4.0, end=6.0, event_type="music") + + writer.write("test.smi") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check for sync blocks + assert 'sync start="1000"' in written_content + assert 'sync start="4000"' in written_content + assert "Hello world" in written_content + assert "[music playing]" in written_content + assert "John Doe" in written_content + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_with_long_text(self, mock_makedirs, mock_file): + """Test writing with long text that should be split into lines.""" + writer = SAMICaptionWriter() + long_text = "This is a very long text that should be split into multiple lines because it exceeds the character limit" + writer.add_caption(start=1.0, end=3.0, text=long_text, event_type="speech") + + writer.write("test.smi") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check that line breaks are added + assert "
" in written_content + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_with_html_entities(self, mock_makedirs, mock_file): + """Test writing with text containing HTML entities.""" + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, end=3.0, text="Hello & ", event_type="speech" + ) + + writer.write("test.smi") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check that HTML entities are escaped + assert "&" in written_content + assert "<" in written_content + assert ">" in written_content + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_sorting(self, mock_makedirs, mock_file): + """Test that captions are sorted by start time.""" + writer = SAMICaptionWriter() + writer.add_caption(start=3.0, end=4.0, text="Third", event_type="speech") + writer.add_caption(start=1.0, end=2.0, text="First", event_type="speech") + writer.add_caption(start=2.0, end=3.0, text="Second", event_type="speech") + + writer.write("test.smi") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Find positions of the sync blocks + first_pos = written_content.find('sync start="1000"') + second_pos = written_content.find('sync start="2000"') + third_pos = written_content.find('sync start="3000"') + + assert first_pos < second_pos < third_pos + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_custom_language_and_title(self, mock_makedirs, mock_file): + """Test writing with custom language and title.""" + writer = SAMICaptionWriter(language="fr-FR", title="Test Title") + writer.add_caption(start=1.0, end=3.0, text="Bonjour", event_type="speech") + + writer.write("test.smi") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check for custom title + assert "Test Title" in written_content + + # Check for custom language class + assert "FRFR" in written_content # fr-FR becomes FRFR + assert "lang: fr-FR" in written_content + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_clearing_syncs(self, mock_makedirs, mock_file): + """Test that clearing syncs are added appropriately.""" + writer = SAMICaptionWriter() + writer.add_caption(start=1.0, end=2.0, text="First", event_type="speech") + writer.add_caption( + start=4.0, end=5.0, text="Second", event_type="speech" + ) # Gap > 500ms + + writer.write("test.smi") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check for clearing syncs + assert 'sync start="2000"' in written_content # End of first caption + assert 'sync start="5000"' in written_content # End of second caption + assert " " in written_content # Non-breaking space for clearing + + def test_write_creates_directory(self): + """Test that write creates the output directory if it doesn't exist.""" + with tempfile.TemporaryDirectory() as temp_dir: + output_path = os.path.join(temp_dir, "subdir", "test.smi") + writer = SAMICaptionWriter() + writer.add_caption(start=1.0, end=2.0, text="Test", event_type="speech") + + writer.write(output_path) + + # Check that the directory was created + assert os.path.exists(os.path.dirname(output_path)) + assert os.path.exists(output_path) + + def test_write_real_file(self): + """Test writing to a real file.""" + with tempfile.NamedTemporaryFile( + mode="w", suffix=".smi", delete=False + ) as temp_file: + temp_path = temp_file.name + + try: + writer = SAMICaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, + speaker="John Doe", + text="Hello world", + event_type="speech", + ) + + writer.write(temp_path) + + # Read the file and check content + with open(temp_path, "r", encoding="utf-8") as f: + content = f.read() + + assert "" in content + assert "" in content + assert "Hello world" in content + assert "John Doe" in content + assert 'sync start="1000"' in content + + finally: + os.unlink(temp_path) diff --git a/tests/test_srt_writer.py b/tests/test_srt_writer.py index 184519a..bf99c02 100644 --- a/tests/test_srt_writer.py +++ b/tests/test_srt_writer.py @@ -15,15 +15,15 @@ def test_add_caption_text_and_silence_music(): writer.add_caption(1.0, 1.2, speaker="SPEAKER_00", text="", event_type="silence") # Long silence (should add) writer.add_caption(1.2, 3.0, speaker="SPEAKER_00", text="", event_type="silence") - # Music event (should skip) + # Music event (should add) writer.add_caption(3.0, 4.5, event_type="music") - # Long music event (should add) + # Music event (should add) writer.add_caption(4.5, 6.0, event_type="music") caps = writer._captions # Expect 1 speech, 1 silence, and 1 music caption, 3 entries total - assert len(caps) == 3 - assert caps[0].text == "Hello world." + assert len(caps) == 4 + assert caps[0]["text"] == "Hello world." assert caps[1]["text"] == "[SILENCE]" assert caps[2]["text"] == "[MUSIC PLAYING]" diff --git a/tests/test_transcriber.py b/tests/test_transcriber.py index eff4aaf..14f999f 100644 --- a/tests/test_transcriber.py +++ b/tests/test_transcriber.py @@ -63,14 +63,13 @@ def test_parse_line(transcriber): assert wt2.is_punctuation is True -@patch("tempfile.NamedTemporaryFile") -@patch("subprocess.run") -@patch("whisper.load_model") -@patch("os.remove") -def test_transcribe_audio_python_api( - mock_remove, mock_load, mock_run, mock_tempfile, transcriber -): +@patch("captionalchemy.tools.captioning.transcriber.tempfile.NamedTemporaryFile") +@patch("captionalchemy.tools.captioning.transcriber.subprocess.run") +@patch("captionalchemy.tools.captioning.transcriber.whisper.load_model") +@patch("captionalchemy.tools.captioning.transcriber.os.remove") +def test_transcribe_audio_python_api(mock_remove, mock_load, mock_run, mock_tempfile): """Test transcribing audio using Whisper Python API.""" + transcriber = Transcriber() # Mock temp file fake = Mock() fake.name = "/tmp/fake.wav" @@ -96,8 +95,8 @@ def test_transcribe_audio_python_api( words = transcriber.transcribe_audio( audio_file="in.wav", - start=1.0, - end=2.0, + start=0.0, + end=1.1, model="base", whisper_build_path=None, whisper_model_path=None, @@ -108,6 +107,9 @@ def test_transcribe_audio_python_api( assert words[0].word == "Hello" assert abs(words[0].start - 0.0) < 1e-3 assert abs(words[0].end - 0.5) < 1e-3 + assert words[1].word == "world" assert abs(words[1].start - 0.6) < 1e-3 assert abs(words[1].end - 1.1) < 1e-3 + + mock_remove.assert_called_once_with(fake.name) diff --git a/tests/test_vtt_writer.py b/tests/test_vtt_writer.py new file mode 100644 index 0000000..bc5f96b --- /dev/null +++ b/tests/test_vtt_writer.py @@ -0,0 +1,442 @@ +import tempfile +import os +from unittest.mock import patch, mock_open + +from captionalchemy.tools.captioning.writers.vtt_writer import VTTCaptionWriter + + +class TestVTTCaptionWriter: + """Test suite for VTTCaptionWriter class.""" + + def test_init_default_values(self): + """Test initialization with default values.""" + writer = VTTCaptionWriter() + + assert writer.min_silence_duration == 1.0 + assert writer.min_music_duration == 1.5 + assert writer.last_speaker is None + assert writer._captions == [] + + def test_init_custom_values(self): + """Test initialization with custom values.""" + writer = VTTCaptionWriter(min_silence_duration=2.0, min_music_duration=3.0) + + assert writer.min_silence_duration == 2.0 + assert writer.min_music_duration == 3.0 + + def test_add_caption_with_text(self): + """Test adding a caption with text.""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, + speaker="John Doe", + text="Hello world", + event_type="speech", + ) + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["start"] == 1.0 + assert caption["end"] == 3.0 + assert caption["speaker"] == "John Doe" + assert caption["text"] == "Hello world" + assert caption["event_type"] == "speech" + + def test_add_caption_with_newlines(self): + """Test adding a caption with newlines that should be replaced.""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, end=3.0, text="Hello\nworld\ntest", event_type="speech" + ) + + caption = writer._captions[0] + assert caption["text"] == "Hello world test" + + def test_add_caption_music_long_duration(self): + """Test adding a music caption with long duration.""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, # 2 seconds, > min_music_duration (1.5) + event_type="music", + ) + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["text"] == "[MUSIC PLAYING]" + + def test_add_caption_music_short_duration(self): + """Test adding a music caption with short duration (should be skipped).""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=2.0, # 1 second, < min_music_duration (1.5) + event_type="music", + ) + + assert len(writer._captions) == 0 + + def test_add_caption_silence_long_duration(self): + """Test adding a silence caption with long duration.""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, # 2 seconds, > min_silence_duration (1.0) + event_type="silence", + ) + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["text"] == "[SILENCE]" + + def test_add_caption_silence_short_duration(self): + """Test adding a silence caption with short duration (should be skipped).""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=1.5, # 0.5 seconds, < min_silence_duration (1.0) + event_type="silence", + ) + + assert len(writer._captions) == 0 + + def test_add_caption_with_label(self): + """Test adding a caption with a label.""" + writer = VTTCaptionWriter() + writer.add_caption(start=1.0, end=3.0, event_type="other", label="applause") + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["text"] == "[APPLAUSE]" + + def test_add_caption_fallback(self): + """Test adding a caption with fallback text.""" + writer = VTTCaptionWriter() + writer.add_caption(start=1.0, end=3.0, event_type="unknown") + + assert len(writer._captions) == 1 + caption = writer._captions[0] + assert caption["text"] == "[AUDIO]" + + def test_format_timestamp_zero(self): + """Test timestamp formatting for zero seconds.""" + result = VTTCaptionWriter._format_timestamp(0.0) + assert result == "00:00:00.000" + + def test_format_timestamp_basic(self): + """Test basic timestamp formatting.""" + result = VTTCaptionWriter._format_timestamp(65.5) + assert result == "00:01:05.500" + + def test_format_timestamp_with_hours(self): + """Test timestamp formatting with hours.""" + result = VTTCaptionWriter._format_timestamp(3665.123) + assert result == "01:01:05.123" + + def test_format_timestamp_edge_cases(self): + """Test timestamp formatting edge cases.""" + # Test milliseconds rounding + result = VTTCaptionWriter._format_timestamp(1.9999) + assert result == "00:00:01.999" + + # Test very small time + result = VTTCaptionWriter._format_timestamp(0.001) + assert result == "00:00:00.001" + + # Test large time + result = VTTCaptionWriter._format_timestamp(7323.456) + assert result == "02:02:03.456" + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_empty_captions(self, mock_makedirs, mock_file): + """Test writing with no captions.""" + writer = VTTCaptionWriter() + writer.write("test.vtt") + + mock_makedirs.assert_called_once_with(".", exist_ok=True) + mock_file.assert_called_once_with("test.vtt", "w", encoding="utf-8") + + # Check that the WebVTT header is written + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + assert written_content == "WEBVTT\n\n" + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_with_speech_captions(self, mock_makedirs, mock_file): + """Test writing with speech captions.""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, + speaker="John Doe", + text="Hello world", + event_type="speech", + ) + writer.add_caption( + start=4.0, + end=6.0, + speaker="Jane Smith", + text="Hi there", + event_type="speech", + ) + + writer.write("test.vtt") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check for WebVTT header + assert "WEBVTT" in written_content + + # Check for timestamps + assert "00:00:01.000 --> 00:00:03.000" in written_content + assert "00:00:04.000 --> 00:00:06.000" in written_content + + # Check for voice tags + assert "Hello world" in written_content + assert "Hi there" in written_content + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_with_same_speaker_consecutive(self, mock_makedirs, mock_file): + """Test writing with consecutive captions from the same speaker.""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, + speaker="John Doe", + text="Hello world", + event_type="speech", + ) + writer.add_caption( + start=4.0, + end=6.0, + speaker="John Doe", + text="How are you?", + event_type="speech", + ) + + writer.write("test.vtt") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check that voice tags are still used for consistency + assert "Hello world" in written_content + assert "How are you?" in written_content + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_with_non_speech_captions(self, mock_makedirs, mock_file): + """Test writing with non-speech captions.""" + writer = VTTCaptionWriter() + writer.add_caption(start=1.0, end=3.0, event_type="music") + writer.add_caption(start=4.0, end=6.0, event_type="silence") + + writer.write("test.vtt") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check for timestamps + assert "00:00:01.000 --> 00:00:03.000" in written_content + assert "00:00:04.000 --> 00:00:06.000" in written_content + + # Check for event descriptions (no voice tags) + assert "[MUSIC PLAYING]" in written_content + assert "[SILENCE]" in written_content + + # Ensure no voice tags are used for non-speech + assert " 00:00:02.000") + second_pos = written_content.find("00:00:02.000 --> 00:00:03.000") + third_pos = written_content.find("00:00:03.000 --> 00:00:04.000") + + assert first_pos < second_pos < third_pos + + @patch("builtins.open", new_callable=mock_open) + @patch("os.makedirs") + def test_write_mixed_content(self, mock_makedirs, mock_file): + """Test writing with mixed speech and non-speech content.""" + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, end=3.0, speaker="John Doe", text="Hello", event_type="speech" + ) + writer.add_caption(start=4.0, end=6.0, event_type="music") + writer.add_caption( + start=7.0, + end=9.0, + speaker="Jane Smith", + text="Goodbye", + event_type="speech", + ) + + writer.write("test.vtt") + + handle = mock_file.return_value + written_content = "".join(call[0][0] for call in handle.write.call_args_list) + + # Check for mixed content + assert "Hello" in written_content + assert "[MUSIC PLAYING]" in written_content + assert "Goodbye" in written_content + + def test_write_creates_directory(self): + """Test that write creates the output directory if it doesn't exist.""" + with tempfile.TemporaryDirectory() as temp_dir: + output_path = os.path.join(temp_dir, "subdir", "test.vtt") + writer = VTTCaptionWriter() + writer.add_caption(start=1.0, end=2.0, text="Test", event_type="speech") + + writer.write(output_path) + + # Check that the directory was created + assert os.path.exists(os.path.dirname(output_path)) + assert os.path.exists(output_path) + + def test_write_real_file(self): + """Test writing to a real file.""" + with tempfile.NamedTemporaryFile( + mode="w", suffix=".vtt", delete=False + ) as temp_file: + temp_path = temp_file.name + + try: + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, + speaker="John Doe", + text="Hello world", + event_type="speech", + ) + writer.add_caption(start=4.0, end=6.0, event_type="music") + + writer.write(temp_path) + + # Read the file and check content + with open(temp_path, "r", encoding="utf-8") as f: + content = f.read() + + assert "WEBVTT" in content + assert "00:00:01.000 --> 00:00:03.000" in content + assert "00:00:04.000 --> 00:00:06.000" in content + assert "Hello world" in content + assert "[MUSIC PLAYING]" in content + + finally: + os.unlink(temp_path) + + def test_write_with_special_characters(self): + """Test writing with special characters in text.""" + with tempfile.NamedTemporaryFile( + mode="w", suffix=".vtt", delete=False + ) as temp_file: + temp_path = temp_file.name + + try: + writer = VTTCaptionWriter() + writer.add_caption( + start=1.0, + end=3.0, + speaker="José María", + text="Hola, ¿cómo estás?", + event_type="speech", + ) + + writer.write(temp_path) + + # Read the file and check content + with open(temp_path, "r", encoding="utf-8") as f: + content = f.read() + + assert "José María" in content + assert "Hola, ¿cómo estás?" in content + + finally: + os.unlink(temp_path) + + def test_write_with_long_timestamps(self): + """Test writing with long timestamps (hours).""" + with tempfile.NamedTemporaryFile( + mode="w", suffix=".vtt", delete=False + ) as temp_file: + temp_path = temp_file.name + + try: + writer = VTTCaptionWriter() + writer.add_caption( + start=3661.5, # 1 hour, 1 minute, 1.5 seconds + end=3665.0, # 1 hour, 1 minute, 5 seconds + text="Long time", + event_type="speech", + ) + + writer.write(temp_path) + + # Read the file and check content + with open(temp_path, "r", encoding="utf-8") as f: + content = f.read() + + assert "01:01:01.500 --> 01:01:05.000" in content + + finally: + os.unlink(temp_path) + + def test_multiple_captions_same_time(self): + """Test handling multiple captions with the same start time.""" + writer = VTTCaptionWriter() + writer.add_caption(start=1.0, end=2.0, text="First", event_type="speech") + writer.add_caption(start=1.0, end=2.0, text="Second", event_type="speech") + + assert len(writer._captions) == 2 + + # Both should be preserved in the output + with tempfile.NamedTemporaryFile( + mode="w", suffix=".vtt", delete=False + ) as temp_file: + temp_path = temp_file.name + + try: + writer.write(temp_path) + + with open(temp_path, "r", encoding="utf-8") as f: + content = f.read() + + assert "First" in content + assert "Second" in content + + finally: + os.unlink(temp_path) + + def test_custom_duration_thresholds(self): + """Test custom duration thresholds for silence and music.""" + writer = VTTCaptionWriter(min_silence_duration=0.5, min_music_duration=1.0) + + # These should now be included with the lower thresholds + writer.add_caption(start=1.0, end=1.6, event_type="silence") # 0.6 seconds + writer.add_caption(start=2.0, end=3.1, event_type="music") # 1.1 seconds + + assert len(writer._captions) == 2 + assert writer._captions[0]["text"] == "[SILENCE]" + assert writer._captions[1]["text"] == "[MUSIC PLAYING]" From 9a295e2e1176fabe9485b89166f6c31701bbd588 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Thu, 3 Jul 2025 10:10:09 -0400 Subject: [PATCH 18/21] fixing linter check --- tests/test_sami_writer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_sami_writer.py b/tests/test_sami_writer.py index 0d8fe27..cbec808 100644 --- a/tests/test_sami_writer.py +++ b/tests/test_sami_writer.py @@ -194,7 +194,8 @@ def test_write_with_captions(self, mock_makedirs, mock_file): def test_write_with_long_text(self, mock_makedirs, mock_file): """Test writing with long text that should be split into lines.""" writer = SAMICaptionWriter() - long_text = "This is a very long text that should be split into multiple lines because it exceeds the character limit" + long_text = """This is a very long text that should be split into + multiple lines because it exceeds the character limit""" writer.add_caption(start=1.0, end=3.0, text=long_text, event_type="speech") writer.write("test.smi") From 1712446ec9da93e617bef4d7f3eb1ef5860e7945 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Thu, 3 Jul 2025 10:16:23 -0400 Subject: [PATCH 19/21] formatting --- .gitignore | 1 - tests/test_main.py | 76 ++++++++++++++++++++-------------------------- 2 files changed, 33 insertions(+), 44 deletions(-) diff --git a/.gitignore b/.gitignore index 9607835..0f325eb 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,3 @@ captionalchemy.egg-info # Whisper whisper.cpp /whisper.cpp - diff --git a/tests/test_main.py b/tests/test_main.py index cc83766..93289f5 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -5,11 +5,6 @@ import shutil import logging -# from captionalchemy.tools.audio_analysis.audio_segment_integration import ( -# AudioEvent, -# EventType, -# ) - from captionalchemy.caption import run_pipeline, _build_arg_parser, main @@ -396,16 +391,14 @@ def test_pipeline_with_existing_video_file(self, tmp_path): video_file = tmp_path / "test_video.mp4" video_file.write_text("dummy video content") - with patch("captionalchemy.caption.embed_faces"), patch( - "captionalchemy.caption.VideoManager" - ) as mock_vm_class, patch("captionalchemy.caption.extract_audio"), patch( - "captionalchemy.caption.get_speech_segments", return_value=[] - ), patch( - "captionalchemy.caption.detect_non_speech_segments", return_value=[] - ), patch( - "captionalchemy.caption.torch.cuda.is_available", return_value=False - ), patch( - "captionalchemy.caption.whisper.load_model" + with ( + patch("captionalchemy.caption.embed_faces"), + patch("captionalchemy.caption.VideoManager") as mock_vm_class, + patch("captionalchemy.caption.extract_audio"), + patch("captionalchemy.caption.get_speech_segments", return_value=[]), + patch("captionalchemy.caption.detect_non_speech_segments", return_value=[]), + patch("captionalchemy.caption.torch.cuda.is_available", return_value=False), + patch("captionalchemy.caption.whisper.load_model"), ): mock_vm = Mock() @@ -691,9 +684,10 @@ def test_main_with_verbose_logging( mock_logger = Mock() mock_get_logger.return_value = mock_logger - with patch("sys.argv", ["captionalchemy", "video.mp4", "--verbose"]), patch( - "captionalchemy.caption.logging.getLogger" - ) as mock_root_logger: + with ( + patch("sys.argv", ["captionalchemy", "video.mp4", "--verbose"]), + patch("captionalchemy.caption.logging.getLogger") as mock_root_logger, + ): main() @@ -745,9 +739,10 @@ def test_main_info_logging_default( ): """Test that main sets INFO logging by default.""" - with patch("sys.argv", ["captionalchemy", "video.mp4"]), patch( - "captionalchemy.caption.logging.getLogger" - ) as mock_root_logger: + with ( + patch("sys.argv", ["captionalchemy", "video.mp4"]), + patch("captionalchemy.caption.logging.getLogger") as mock_root_logger, + ): main() @@ -762,16 +757,14 @@ class TestIntegration: def test_pipeline_environment_variables(self): """Test that environment variables are properly used.""" - with patch("captionalchemy.caption.embed_faces"), patch( - "captionalchemy.caption.VideoManager" - ), patch("captionalchemy.caption.extract_audio"), patch( - "captionalchemy.caption.get_speech_segments" - ) as mock_vad, patch( - "captionalchemy.caption.detect_non_speech_segments", return_value=[] - ), patch( - "captionalchemy.caption.torch.cuda.is_available", return_value=False - ), patch( - "captionalchemy.caption.whisper.load_model" + with ( + patch("captionalchemy.caption.embed_faces"), + patch("captionalchemy.caption.VideoManager"), + patch("captionalchemy.caption.extract_audio"), + patch("captionalchemy.caption.get_speech_segments") as mock_vad, + patch("captionalchemy.caption.detect_non_speech_segments", return_value=[]), + patch("captionalchemy.caption.torch.cuda.is_available", return_value=False), + patch("captionalchemy.caption.whisper.load_model"), ): mock_vad.return_value = [] @@ -802,18 +795,15 @@ def __exit__(self, *args): if os.path.exists(self.name): shutil.rmtree(self.name) - with patch("tempfile.TemporaryDirectory", MockTempDir), patch( - "captionalchemy.caption.embed_faces" - ), patch("captionalchemy.caption.VideoManager"), patch( - "captionalchemy.caption.extract_audio" - ), patch( - "captionalchemy.caption.get_speech_segments", return_value=[] - ), patch( - "captionalchemy.caption.detect_non_speech_segments", return_value=[] - ), patch( - "captionalchemy.caption.torch.cuda.is_available", return_value=False - ), patch( - "captionalchemy.caption.whisper.load_model" + with ( + patch("tempfile.TemporaryDirectory", MockTempDir), + patch("captionalchemy.caption.embed_faces"), + patch("captionalchemy.caption.VideoManager"), + patch("captionalchemy.caption.extract_audio"), + patch("captionalchemy.caption.get_speech_segments", return_value=[]), + patch("captionalchemy.caption.detect_non_speech_segments", return_value=[]), + patch("captionalchemy.caption.torch.cuda.is_available", return_value=False), + patch("captionalchemy.caption.whisper.load_model"), ): run_pipeline("test_video.mp4") From 779c424047627f704930b25ac9ce06999cd3c8c1 Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Thu, 3 Jul 2025 10:23:55 -0400 Subject: [PATCH 20/21] adding pytest-cov to dev dependencies --- .gitignore | 4 ++++ pyproject.toml | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 0f325eb..688c13d 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,7 @@ captionalchemy.egg-info # Whisper whisper.cpp /whisper.cpp + +# Coverage +.coverage +.coverage.* diff --git a/pyproject.toml b/pyproject.toml index 6bd3817..33bf528 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ requires-python = ">=3.10" [project.optional-dependencies] -dev = ["pytest==8.4.0", "mypy==1.16.0", "flake8==7.2.0", "black==25.1.0"] +dev = ["pytest==8.4.0", "mypy==1.16.0", "flake8==7.2.0", "black==25.1.0", "pytest-cov==6.2.1"] [tool.setuptools.packages.find] where = ["src"] From b88872ed256fb4b0c248689f348d77e0f15cc05b Mon Sep 17 00:00:00 2001 From: Ben Batman Date: Thu, 3 Jul 2025 10:40:00 -0400 Subject: [PATCH 21/21] installing wget on windows runner --- .coverage | Bin 53248 -> 0 bytes .github/workflows/ci.yml | 2 ++ .gitignore | 1 + 3 files changed, 3 insertions(+) delete mode 100644 .coverage diff --git a/.coverage b/.coverage deleted file mode 100644 index e575f9271e600d8959c2a0ae0e881b1272dba5ac..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 53248 zcmeI4eQX>@6~K4z*0;C!HFlG5O+)cYsCC@fKCDWG5ZWY#sH$i&NgOH?CELBbvAxaR z?qzq+iBp>R5~K*Es)9d2mC{gEMHN&6{((XTq^JdnA~`J;5{Qq45492bXp7RMA;#go z{qWs6c2vZgD&dWMGrO}h^XC2L&CJcMXEtxV&v3aqV%b&QRc{wM1yK~%sj46d3VgEg z37rfVQg{QQ*mAwlr6LTN&u7`^giPX5f$hi+vHr~A?5*inGgalw=_9hPl;8n5B0vO) z01^2ACeYZJk@L%!iw7QY^BoCLHkc__bJRlzUUX2luDs{E{W&A@1p< z+a>r*w`l_rO;zq4fr3hwUB=oMZQU%5ai>q!XEiKYP|rQCudQuW3$CSB4AWE7bPU(9 zOqK89rMk<@^D5Yf&_W0BTD5_)9z=y6^m7KXTIe)0&yzDq9%LSY%9&P(wV4q`LuRke zX1!Vjo*G*R=v=-*el{)0Rzr+nO^(jxFlvs+-2HI-hO)KDC*% z&r`#+`czZ`EZ~-OBENcxD1cS+JyQm>?#iGP_o9kt&j47;Hf_>yx)cdrZcD{ z@*BDrGM!ly&<1BOQ;p>bIlp3sctG*YC}^!9klRcsY{_|s(eP1~>rUMc2a>A|l3|YK zC>i#V;0dVHFDLTrS1hDXkWL#6E|LvdlJm=!i4DI;;JAdJjy5GQrr8DNOa$+}`Ze7h zQ?IN$MnlDWIuiLc%ND9QRzVv?KXGtd%&YRD(2_4NyA-iPcVRj-#Ey7y--$1z}9#c9j= zco?hpjOh-JLQaobw5$r(O@D5H9QQgP({NO;4p+`fYnx>_6IjORc#l`7aB{(6ei%Ak z*uUVSTGh>RFkj&`?3W5-KXy)(;}0T`QSE_F=|C5qQ)XCb@DxX^u>64Qz0gc5oYOGm<$EoAxpU11@~qX|)q%*Mb9YeF-?Y`7K!V zDTJwl?>eboT;PlBk^q0?hyW2F0z`la5CI}U1c(3;AOb{y2oQnmk3d|Ei88MLW9-iY zdlNizM1Tko0U|&IhyW2F0z`la5CI}U1c<=3Ngxx8-^{{)@%RW6<16yPKLLDnVD-RY z3bcx`lL9-*-nljvq*@aJB0vO)01+SpM1Tko0U|&IhyW2F0wIA+{AMxy6(E+1uTX*? z0r35QSMptfeVeV#y_oxa_Eh$f?6S-YnN8`_>HXWS1U<>$(0APG4lKm>>Y5g-CY zfCyaE1ctjLq0wL1;=mr$La|uXUD%jf*j%!Wn(GwG#;DP6n5NZV(rdU)yx-ddYS{(H zE){|}y@I$^!t%^iiV2)7ehzx~pjHFP%8v`f~ z2c49)ohB&IM3smb)tuo1Ou;tx!a+Eei4D^cv9y-{3tQ)*8~X#*UT%Z5F$tt?iISl| z660tAW^Fn!RKw9RAZc$LBuzvpX%0pI*5RX7Zn~PcbKVY_8ZnUlz~af)K%U`xbIY&@ zqN19yW=qBQ|6R$80(&-jG4nIFF85M)pE9KUB6~WwHPx49sUz7Bq|anFrzfxJ>7FW2 z1c(3;AOb{y2oM25*djIh#Yn!mwXFZU#-)Z8O{xFM`oEJ)jr$i-tiS$eyQBt>QcBDE zVgB`ht|m1~i!8Bi{h$3G2y9#bXDU)-Ta+4JZT+8)+DztM|EH=_WBX#MbjA8#8I~HW zqLJ2I|0hSK#>Qw0z3Td3jvCdRm23O@KM}RGmi~L?`oF`F8e1Y|xX}7vdRS^qL@8+w z#jCCV;}1!V2NqFw%lbbCmfo*KG-J)K|LOcc`Z7cWhyW2F0z`la5CI}U1c(3;AOb|- zx+V}8<3bkS|BGx|fIo6XfCvx)B0vO)01+SpM1Tko0U|&Ih`@D7ARf;u`2N2$cTiv# z+1uM2UtfFd!r`Yc{OQ+g^7kA+ z(6=-TDQ?S-oc!)u;S0b2`(u;CJwrF`Uz)+#j32x8#AB1^PAh+9avFhlzCm`a+#MxV?-kw=D)X^a-jSqFq{Nt!F ztqC)yFV8$D2#MMYNkK(vN%X*-fTgz}tP@1pi@MW`1TP045pQ1UMql{YfCQ!QkY={e zT=<4KnOG{Uo4oUxB@Z`VUVY1@q06E$b$Nf@%kb7)|JpiT5^st_`a9xNZ@hb9weXW2 ze>nWa*Sm#Ro_god^jEqQF?g1Xi5;S(K#1@EJK0wSc8Q&1XV~lPIQuPomA%YfWItx# zXWwDZuqWB$kboQ!AOb{y2oM1xKm>>Y5g-CYfCvx)A`l?~{3Su`?DPW0ydanJf^60c zG8r#Od%rn|sg!r8C|;0EdVwr^K_cM=9UWdENnQ|-dqFHF!4Cl5`+xiuf*cVb0z`la z5CI}U1c(3;AOb{y2oM1xaQza%_y4s1zkb`7G$R5;fCvx)B0vO)01+SpM1Tko0V03| zaQ}Z7dlvuy|2&)kaGJfr{=)tUdjMWzQ|y=QC3ciO5BmVV2T8~g0U|&IhyW2F0z`la z5CI}U1c(3;AOed|Al?P%^~XEWW9a43%c7S-FO6ObJq5iadNO(m^g7U!(2JuNLl4jY F|2Il>C0hUh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a3a0b1d..98ac84c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -65,6 +65,8 @@ jobs: choco install -y ffmpeg # Add ffmpeg to PATH echo "C:\ProgramData\chocolatey\lib\ffmpeg\tools\ffmpeg\bin" >> $GITHUB_PATH + # Install wget + choco install -y wget pip3 install torch torchvision torchaudio - name: Upgrade pip and dependencies diff --git a/.gitignore b/.gitignore index 688c13d..3b420f9 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ venv __pycache__ .mypy_cache captionalchemy.egg-info +.pytest_cache *.env .env