highlight.py (4731B)
1 import sys 2 import json 3 import re 4 import pathlib 5 from datetime import timedelta 6 7 # ANSI escape codes for terminal colors 8 class Colors: 9 GREEN = '\033[92m' # Highlight for valid ad text 10 RED = '\033[91m' # Highlight for invalid target ranges 11 BLUE = '\033[94m' # Header color 12 RESET = '\033[0m' # Reset to default color 13 14 def timestamp_to_seconds(ts_str: str) -> float: 15 """Converts HH:MM:SS.mmm string to total seconds. Returns -1.0 on failure.""" 16 try: 17 parts = ts_str.split(':') 18 h, m = int(parts[0]), int(parts[1]) 19 s_ms = parts[2].split('.') 20 s, ms = int(s_ms[0]), int(s_ms[1]) 21 return h * 3600 + m * 60 + s + ms / 1000.0 22 except (ValueError, IndexError): 23 return -1.0 24 25 def visualize_entry(entry: dict): 26 """Parses and prints a single entry from the jsonl file with highlighting.""" 27 entry_id = entry.get("id", "N/A") 28 full_text = entry.get("text", "") 29 target_str = entry.get("target", "") 30 31 print(f"{Colors.BLUE}{'='*80}\nID: {entry_id}{Colors.RESET}") 32 33 # --- 1. Parse the transcript from the 'text' field --- 34 transcript_lines = [] 35 # Regex to find the start of the actual transcript text 36 transcript_match = re.search(r"Transcript:\n(.*)", full_text, re.DOTALL) 37 if not transcript_match: 38 print(f"{Colors.RED}Could not find transcript text in entry.{Colors.RESET}") 39 return 40 41 raw_transcript = transcript_match.group(1).strip() 42 line_pattern = re.compile(r"\[(\d{2}:\d{2}:\d{2}\.\d{3}) --> (\d{2}:\d{2}:\d{2}\.\d{3})\](.*)") 43 44 min_transcript_time, max_transcript_time = float('inf'), float('-inf') 45 46 for line in raw_transcript.split('\n'): 47 match = line_pattern.match(line) 48 if match: 49 start_str, end_str, text = match.groups() 50 start_sec = timestamp_to_seconds(start_str) 51 end_sec = timestamp_to_seconds(end_str) 52 transcript_lines.append({"start": start_sec, "end": end_sec, "line": line}) 53 min_transcript_time = min(min_transcript_time, start_sec) 54 max_transcript_time = max(max_transcript_time, end_sec) 55 56 # --- 2. Parse and validate the target ranges --- 57 valid_target_ranges = [] 58 has_invalid_target = False 59 60 if target_str: 61 for target_line in target_str.strip().split('\n'): 62 if '-' not in target_line: 63 has_invalid_target = True 64 continue 65 66 start_target_str, end_target_str = target_line.split('-', 1) 67 start_target_sec = timestamp_to_seconds(start_target_str.strip()) 68 end_target_sec = timestamp_to_seconds(end_target_str.strip()) 69 70 # Validation checks 71 is_valid_format = start_target_sec != -1.0 and end_target_sec != -1.0 72 is_valid_range = is_valid_format and start_target_sec < end_target_sec 73 # Check if the target overlaps with the transcript's time frame at all 74 is_overlapping = is_valid_range and start_target_sec < max_transcript_time and end_target_sec > min_transcript_time 75 76 if is_overlapping: 77 valid_target_ranges.append({"start": start_target_sec, "end": end_target_sec}) 78 else: 79 has_invalid_target = True 80 81 # --- 3. Print results with highlighting --- 82 print(f"TARGET: ", end="") 83 if has_invalid_target: 84 print(f"{Colors.RED}{target_str or 'EMPTY'}{Colors.RESET}") 85 else: 86 print(f"{Colors.GREEN}{target_str or 'EMPTY'}{Colors.RESET}") 87 print("-" * 30) 88 89 for line_info in transcript_lines: 90 is_ad_line = False 91 for target_range in valid_target_ranges: 92 # Check for overlap between the line's time and the target's time 93 if line_info['start'] < target_range['end'] and line_info['end'] > target_range['start']: 94 is_ad_line = True 95 break 96 97 if is_ad_line: 98 print(f"{Colors.GREEN}{line_info['line']}{Colors.RESET}") 99 else: 100 print(line_info['line']) 101 102 def main(filepath: str): 103 try: 104 p = pathlib.Path(filepath) 105 with p.open('r', encoding='utf-8') as f: 106 for line in f: 107 try: 108 entry = json.loads(line) 109 visualize_entry(entry) 110 except json.JSONDecodeError: 111 print(f"{Colors.RED}Error: Could not decode JSON line: {line.strip()}{Colors.RESET}") 112 except FileNotFoundError: 113 print(f"{Colors.RED}Error: File not found at {filepath}{Colors.RESET}") 114 115 if __name__ == "__main__": 116 if len(sys.argv) < 2: 117 print("Usage: python highlight.py <path_to_labels.jsonl>") 118 sys.exit(1) 119 120 main(sys.argv[1])