"add wordlevel highlighting feature (wip)"

main.py

@@ -16,12 +16,12 @@ from fastapi.security import HTTPBasic
 from pydantic import BaseModel, field_validator
 from cachetools import TTLCache
 
-## THIS IS A BREAKING CHANGE. SRT FILE INPUT DEPRECATED. WIP.
-## DONE: separate transcriber from subtitler logic. WIP.
-## DONE: improve loading spinner. WIP (with redirect)
+## THIS IS A BREAKING CHANGE. SRT FILE INPUT DEPRECATED.
+## DONE: separate transcriber from subtitler logic.
+## DONE: improve loading spinner. (redirect)
 ## DONE: fix tempdir cleanup
 ## DONE: add transcription preview component + allow for interactive validation of transcription in-browser.
-## TODO: add word level highlighting option
+## TODO: add word level highlighting option. WIP (word background margins need to be addressed; mobile mode needs work in json mode)
 ## TODO: improve UI
 
 app = FastAPI()
@@ -79,10 +79,14 @@ async def transcribe_api(video_file: MP4Video = Depends(),
         with open(video_path, 'wb') as f:
             shutil.copyfileobj(video_file.file, f)
 
-        transcription = transcriber(video_path, max_words_per_line, task, model_version)
+        transcription_text, transcription_json = transcriber(video_path, max_words_per_line, task, model_version)
 
         uid = str(uuid4())
-        cache[uid] = {"video_path": video_path, "transcription": transcription, "temp_dir_path": temp_dir.name}
+        cache[uid] = {
+            "video_path": video_path,
+            "transcription_text": transcription_text,
+            "transcription_json": transcription_json,
+            "temp_dir_path": temp_dir.name}
         return RedirectResponse(url=f"/process_settings/?uid={uid}", status_code=303)
     
     except Exception as e:
@@ -95,7 +99,8 @@ async def process_settings(request: Request, uid: str):
         raise HTTPException(404, "Data not found")
     return templates.TemplateResponse("process_settings.html", {
         "request": request,
-        "transcription": data["transcription"],
+        "transcription_text": data["transcription_text"],
+        "transcription_json": data["transcription_json"],
         "video_path": data["video_path"],
         "temp_dir_path": data["temp_dir_path"]
     })
@@ -104,15 +109,17 @@ async def process_settings(request: Request, uid: str):
 async def process_video_api(video_path: str = Form(...),
                             temp_dir_path: str = Form(...),
                             srt_string: str = Form(...),
+                            srt_json: str = Form(...),
                             fontsize: Optional[int] = Form(42),
                             font: Optional[str] = Form("Helvetica"),
                             bg_color: Optional[str] = Form("#070a13b3"),
                             text_color: Optional[str] = Form("white"),
+                            highlight_mode: Optional[bool] = Form(False),
                             caption_mode: Optional[str] = Form("desktop"),
                             temp_dir: TemporaryDirectory = Depends(get_temp_dir)
                             ):
     try:
-        output_path = process_video(video_path, srt_string, fontsize, font, bg_color, text_color, caption_mode)
+        output_path = process_video(video_path, srt_string, srt_json, fontsize, font, bg_color, text_color, highlight_mode, caption_mode)
         with open(os.path.join(temp_dir.name, f"{video_path.split('.')[0]}.srt"), 'w+') as temp_srt_file:
             logging.info("Processing the video...")
             temp_srt_file.write(srt_string)    

static/process_settings.html

@@ -8,14 +8,44 @@
         label, select, input, textarea { display: block; width: 100%; margin-bottom: 1rem; }
         textarea { height: 200px; font-family: monospace; }
         input[type="submit"] { background: #4CAF50; color: white; padding: 0.8rem; border: none; cursor: pointer; }
+        .radio-container {
+        display: flex;
+        gap: 1rem;
+        margin-bottom: 1rem;
+        }
+
+        .radio-option {
+        display: flex;
+        flex-direction: column;
+        align-items: center;
+        }
     </style>
 </head>
 <body>
     <form action="/process_video/" method="post">
         <h2>Step 2: Edit Transcription & Style</h2>
+        
+        <div class="radio-container">
+            <div class="radio-option">
+              <label for="mode_normal">Normal</label>
+              <input type="radio" name="highlight_mode" value="false" id="mode_normal" checked onchange="toggleTranscriptionFields()">
+            </div>
+            <div class="radio-option">
+              <label for="mode_highlight">Word-level</label>
+              <input type="radio" name="highlight_mode" value="true" id="mode_highlight" onchange="toggleTranscriptionFields()">
+            </div>
+          </div>
+
+        <!-- Textareas -->
+        <div id="normal_input">
+        <label for="srt_string">Transcription (SRT)</label>
+        <textarea name="srt_string" id="srt_string">{{ transcription_text }}</textarea>
+        </div>
 
-        <label for="srt_string">Correct Transcription</label>
-        <textarea name="srt_string" id="srt_string">{{ transcription }}</textarea>
+        <div id="highlight_input" style="display: none;">
+        <label for="srt_json">Transcription (JSON)</label>
+        <textarea name="srt_json" id="srt_json">{{ transcription_json }}</textarea>
+        </div>
 
         <label for="fontsize">Font size</label>
         <input type="number" name="fontsize" value="42">
@@ -78,6 +108,11 @@
     // Populate dropdowns with defaults
     populateDropdown('font', '/static/fonts.txt', DEFAULT_FONT);
     populateDropdown('text_color', '/static/colors.txt', DEFAULT_COLOR);
+    function toggleTranscriptionFields() {
+        const isHighlight = document.getElementById('mode_highlight').checked;
+        document.getElementById('normal_input').style.display = isHighlight ? 'none' : 'block';
+        document.getElementById('highlight_input').style.display = isHighlight ? 'block' : 'none';
+    }
 </script>
 </body>
 </html>

utils/process_video.py

@@ -1,17 +1,19 @@
-import logging, os
+import logging, os, json
 from utils.subtitler import subtitler
 
 def process_video(invideo_file: str,
                   srt_string:str,
+                  srt_json: str,
                   fontsize:str,
                   font:str,
                   bg_color:str,
                   text_color:str,
+                  highlight_mode: bool,
                   caption_mode:str
                   ):
     invideo_path_parts = os.path.normpath(invideo_file).split(os.path.sep)
     VIDEO_NAME = os.path.basename(invideo_file)
     OUTVIDEO_PATH = os.path.join(os.path.normpath('/'.join(invideo_path_parts[:-1])), f"result_{VIDEO_NAME}")
     logging.info("Subtitling...")
-    subtitler(invideo_file, srt_string, OUTVIDEO_PATH, fontsize, font, bg_color, text_color, caption_mode)
+    subtitler(invideo_file, srt_string, srt_json, OUTVIDEO_PATH, fontsize, font, bg_color, text_color, highlight_mode, caption_mode)
     return OUTVIDEO_PATH

utils/subtitler.py

@@ -1,5 +1,5 @@
 from moviepy.editor import VideoFileClip, CompositeVideoClip, TextClip
-import os
+import os, json
 
 def parse_srt(srt_string):
     """Parse the SRT string and return a list of (start, end, text) for each subtitle."""
@@ -27,31 +27,77 @@ def filter_caption_width(caption_mode:str):
         caption_height_ratio = 0.7
     return caption_width_ratio, caption_height_ratio
 
-def subtitler(video_file:str,
-            srt_string:str,
-            output_file:str,
-            fontsize:int,
+
+def subtitler(video_file: str,
+            srt_string: str,
+            srt_json: str,
+            output_file: str,
+            fontsize: int,
             font: str,
-            bg_color:str,
-            text_color:str,
-            caption_mode:str
+            bg_color: str,
+            text_color: str,
+            highlight_mode: bool,
+            caption_mode: str
             ):
-    """Add subtitles from an SRT string to a video."""
+    """Add subtitles to a video, with optional word-level highlighting."""
     video_file = os.path.abspath(video_file)
     output_file = os.path.abspath(output_file)
     clip = VideoFileClip(filename=video_file, target_resolution=None)
-    subtitles = parse_srt(srt_string)
+
     subtitle_clips = []
+
     caption_width_ratio, caption_height_ratio = filter_caption_width(caption_mode)
+    subtitle_y_position = clip.h * caption_height_ratio
+    if highlight_mode:
+        srt_data = json.loads(json.dumps(eval(srt_json)))
+        for line in srt_data.get("lines", []):
+            line_start = float(line["start"])
+            line_end = float(line["end"])
+            line_text = line["text"]
+
+            base_clip = TextClip(line_text, fontsize=fontsize, color=text_color, font=font, method='label')
+            base_clip = base_clip.set_start(line_start).set_end(line_end)
+
+            # Center the full line
+            line_width = base_clip.w
+            x_center = (clip.w - line_width) // 2
+            base_clip = base_clip.set_position((x_center, subtitle_y_position))
+            subtitle_clips.append(base_clip)
+
+            # Calculate word-level highlight positions
+            current_x = x_center
+            for word_info in line["words"]:
+                word = word_info["word"]
+                word_start = float(word_info["start"])
+                word_end = float(word_info["end"])
+
+                # Create a background-only word clip
+                word_clip = TextClip(word, fontsize=fontsize, color=text_color, stroke_color=text_color, stroke_width=2, font=font,
+                        method='label', bg_color="LightBlue")
+                word_clip = word_clip.set_start(word_start).set_end(word_end)
+                word_clip = word_clip.set_position((current_x - 7.5, subtitle_y_position))
+                subtitle_clips.append(word_clip)
+
+                space_width = TextClip(" ", fontsize=fontsize, font=font, method='label').w
+                current_x += word_clip.w + space_width
+        video = CompositeVideoClip(size=None, clips=[clip] + subtitle_clips)
+        video.write_videofile(output_file, codec='libx264', audio_codec='aac')
+        return
+    # Normal mode
+    subtitles = parse_srt(srt_string)
+    subtitle_x_position = 'center'
+    subtitle_y_position = clip.h * caption_height_ratio
+    text_position = (subtitle_x_position, subtitle_y_position)
     for start, end, text in subtitles:
-        # Create TextClip with specified styling
-        # To get a list of possible color and font values run: print(TextClip.list("font"), '\n\n', TextClip.list("color"))
-        txt_clip = TextClip(text, fontsize=fontsize, color=text_color, font=font, method='caption',
-                            bg_color=bg_color, align='center', size=(clip.w*caption_width_ratio, None))
-        txt_clip = txt_clip.set_position(('center', 'bottom')).set_duration(clip.duration).set_start(start).set_end(end)
-        subtitle_x_position = 'center'
-        subtitle_y_position = clip.h * caption_height_ratio
-        text_position = (subtitle_x_position, subtitle_y_position)                    
-        subtitle_clips.append(txt_clip.set_position(text_position))
+        txt_clip = TextClip(text,
+                            fontsize=fontsize,
+                            color=text_color,
+                            font=font,
+                            method='caption',
+                            bg_color=bg_color,
+                            align='center',
+                            size=(clip.w * caption_width_ratio, None))
+        txt_clip = txt_clip.set_start(start).set_end(end).set_position(text_position)
+        subtitle_clips.append(txt_clip)
     video = CompositeVideoClip(size=None, clips=[clip] + subtitle_clips)
     video.write_videofile(output_file, codec='libx264', audio_codec='aac')

utils/transcriber.py

@@ -19,4 +19,4 @@ def transcriber(invideo_file:str,
                 model_version=model_version,
                 api_name="/predict"
         )
-        return result[0]
+        return result[0], result[3]