aiavatartest / stream_server.py
Spanicin's picture
Update stream_server.py
cc2a8df verified
raw
history blame
7.16 kB
import os
import subprocess
import pathlib
# from fastapi import HTTPException
import asyncio
# Define base directories
BASE_DIR = pathlib.Path('./videos').resolve()
HLS_DIR = pathlib.Path('./hls_videos').resolve()
HLS_DIR.mkdir(exist_ok=True)
# Keep track of video names added to the queue
video_names = []
segment_counter = 1
total_processed_duration = 0
segment_lock = asyncio.Lock()
def is_valid_path(video_name):
"""
Validates the video path to prevent directory traversal attacks.
Args:
video_name (str): Name of the video file.
Returns:
bool: True if valid, False otherwise.
"""
video_path = (BASE_DIR / video_name).resolve()
return str(video_path).startswith(str(BASE_DIR))
def convert_to_hls(input_file, output_playlist, segment_prefix='segment', segment_duration=10):
"""
Converts an MP4 file to HLS .ts segments, maintaining continuity characteristics across segments.
:param input_file: Path to the input MP4 file.
:param output_playlist: Path to the output .m3u8 playlist file.
:param segment_prefix: Prefix for naming the .ts segments. Default is 'segment'.
:param segment_duration: Duration of each segment in seconds. Default is 10 seconds.
"""
if not os.path.exists(input_file):
raise FileNotFoundError(f"Input file '{input_file}' does not exist.")
# FFmpeg command to convert MP4 to HLS segments
os.chmod(input_file, 0o644) # Change permission to read/write for the user
command = [
'ffmpeg',
'-i', input_file, # Input MP4 file
'-c:v', 'libx264', # Video codec, consistent across all segments
'-c:a', 'aac', # Audio codec, consistent across all segments
'-strict', '-2', # Strict flag for AAC codec
'-flags', '-global_header', # Set global header flag for consistency
'-hls_time', str(segment_duration), # Segment duration in seconds
'-hls_list_size', '0', # Keep all segments in the playlist
'-hls_flags', 'append_list+omit_endlist+program_date_time', # Flags to maintain continuity
'-hls_segment_type', 'mpegts', # Ensure the segment type is TS
'-hls_segment_filename', f'{HLS_DIR}/{segment_prefix}%d.ts', # Naming pattern for the .ts segments
'-force_key_frames', f'expr:gte(t,n_forced*{segment_duration})', # Force keyframes for segment duration consistency
'-avoid_negative_ts', 'make_zero', # Avoid negative timestamps
output_playlist # Output .m3u8 playlist file
]
try:
# Run the FFmpeg command
subprocess.run(command, check=True)
print(f"Successfully converted '{input_file}' to HLS segments with playlist '{output_playlist}'.")
except subprocess.CalledProcessError as e:
print(f"Error during conversion: {e}")
except Exception as e:
print(f"An unexpected error occurred: {e}")
def add_video(video_name, output_path, audio_duration):
convert_to_hls(video_name, output_path, segment_duration=audio_duration)
return {"message": f'"{video_name}" added to the streaming queue.'}
async def concatenate_playlists(video_names, base_dir):
"""
Concatenates multiple HLS playlists into a single playlist with unique segment numbering.
Args:
video_names (list): List of video names added to the queue.
base_dir (str): Base directory where HLS files are stored.
request (Request): FastAPI request object to extract base URL.
"""
concatenated_playlist_path = os.path.join(base_dir, 'master.m3u8')
max_segment_duration = 3 # Since we set hls_time to 3
segment_lines = [] # To store segment lines
# Construct base URL from the incoming request
for video_name in video_names:
video_playlist_path = os.path.join(base_dir, f'{video_name}.m3u8')
if os.path.exists(video_playlist_path):
with open(video_playlist_path, 'r') as infile:
lines = infile.readlines()
for line in lines:
line = line.strip()
if line.startswith('#EXTINF'):
# Append EXTINF line
segment_lines.append(line)
elif line.endswith('.ts'):
segment_file = line
# Update segment URI to include full URL
segment_path = f'{segment_file}'
segment_lines.append(segment_path)
elif line.startswith('#EXT-X-BYTERANGE'):
# Include byte range if present
segment_lines.append(line)
elif line.startswith('#EXT-X-ENDLIST'):
# Do not include this here; we'll add it at the end
continue
elif line.startswith('#EXTM3U') or line.startswith('#EXT-X-VERSION') or line.startswith('#EXT-X-PLAYLIST-TYPE') or line.startswith('#EXT-X-TARGETDURATION') or line.startswith('#EXT-X-MEDIA-SEQUENCE'):
# Skip these tags; they'll be added in the concatenated playlist
continue
else:
# Include any other necessary tags
segment_lines.append(line)
# Write the concatenated playlist
with open(concatenated_playlist_path, 'w') as outfile:
outfile.write('#EXTM3U\n')
outfile.write('#EXT-X-PLAYLIST-TYPE:VOD\n')
outfile.write(f'#EXT-X-TARGETDURATION:{max_segment_duration}\n')
outfile.write('#EXT-X-VERSION:4\n')
outfile.write(f'#EXT-X-MEDIA-SEQUENCE:{1}\n') # Starting from segment number 1
for line in segment_lines:
outfile.write(f'{line}\n')
outfile.write('#EXT-X-ENDLIST\n')
def generate_m3u8(video_duration, output_path,segment_duration=3):
# Initialize playlist content
m3u8_content = "#EXTM3U\n"
m3u8_content += "#EXT-X-PLAYLIST-TYPE:VOD\n"
m3u8_content += f"#EXT-X-TARGETDURATION:{segment_duration}\n"
m3u8_content += "#EXT-X-VERSION:4\n"
m3u8_content += "#EXT-X-MEDIA-SEQUENCE:1\n"
# # Calculate the number of full segments and the remaining duration
# segment_duration = int(segment_duration)
# full_segments = int(video_duration // segment_duration)
# remaining_duration = video_duration % segment_duration
# # Add full segments to the playlist
# for i in range(full_segments):
# m3u8_content += f"#EXTINF:{segment_duration:.6f},\n"
# m3u8_content += f"/live_stream/video_stream{i + 1}.ts\n"
# # Add the remaining segment if there's any leftover duration
# if remaining_duration > 0:
# m3u8_content += f"#EXTINF:{remaining_duration:.6f},\n"
# m3u8_content += f"/live_stream/video_stream{full_segments + 1}.ts\n"
# # End the playlist
# m3u8_content += "#EXT-X-ENDLIST\n"
with open(output_path, "w") as file:
file.write(m3u8_content)
print(f"M3U8 playlist saved to {output_path}")