@@ -52,6 +52,12 @@ def get_open_write_kwargs(
5252 return open_kwargs
5353
5454
55+ def video_stream_bit_depth (stream ) -> int :
56+ if stream is None or stream .format is None or not stream .format .components :
57+ return 8
58+ return max (component .bits for component in stream .format .components )
59+
60+
5561class VideoFromFile (VideoInput ):
5662 """
5763 Class representing video input from a file.
@@ -97,6 +103,13 @@ def get_dimensions(self) -> tuple[int, int]:
97103 return stream .width , stream .height
98104 raise ValueError (f"No video stream found in file '{ self .__file } '" )
99105
106+ def get_bit_depth (self ) -> int :
107+ if isinstance (self .__file , io .BytesIO ):
108+ self .__file .seek (0 ) # Reset the BytesIO object to the beginning
109+ with av .open (self .__file , mode = "r" ) as container :
110+ video_stream = container .streams .video [0 ] if len (container .streams .video ) > 0 else None
111+ return video_stream_bit_depth (video_stream )
112+
100113 def get_duration (self ) -> float :
101114 """
102115 Returns the duration of the video in seconds.
@@ -377,25 +390,32 @@ def save_to(
377390 format : VideoContainer = VideoContainer .AUTO ,
378391 codec : VideoCodec = VideoCodec .AUTO ,
379392 metadata : Optional [dict ] = None ,
393+ bit_depth : int | None = None ,
380394 ):
381395 if isinstance (self .__file , io .BytesIO ):
382396 self .__file .seek (0 ) # Reset the BytesIO object to the beginning
383397 with av .open (self .__file , mode = 'r' ) as container :
384398 container_format = container .format .name
385- video_encoding = container .streams .video [0 ].codec .name if len (container .streams .video ) > 0 else None
399+ video_stream = container .streams .video [0 ] if len (container .streams .video ) > 0 else None
400+ video_encoding = video_stream .codec .name if video_stream is not None else None
401+ source_bit_depth = video_stream_bit_depth (video_stream )
386402 reuse_streams = True
387403 if format != VideoContainer .AUTO and format not in container_format .split ("," ):
388404 reuse_streams = False
389405 if codec != VideoCodec .AUTO and codec != video_encoding and video_encoding is not None :
390406 reuse_streams = False
407+ if bit_depth is not None and video_encoding is not None and bit_depth != source_bit_depth :
408+ reuse_streams = False
391409 if self .__start_time or self .__duration :
392410 reuse_streams = False
393411
394412 if not reuse_streams :
413+ if bit_depth is None :
414+ bit_depth = source_bit_depth
395415 components = self .get_components_internal (container )
396416 video = VideoFromComponents (components )
397417 return video .save_to (
398- path , format = format , codec = codec , metadata = metadata
418+ path , format = format , codec = codec , metadata = metadata , bit_depth = bit_depth ,
399419 )
400420
401421 streams = container .streams
@@ -451,8 +471,10 @@ class VideoFromComponents(VideoInput):
451471 Class representing video input from tensors.
452472 """
453473
454- def __init__ (self , components : VideoComponents ):
474+ def __init__ (self , components : VideoComponents , bit_depth : int = 8 ):
455475 self .__components = components
476+ # Tensor components have no inherent bit depth; this is the depth used when encoding.
477+ self .__bit_depth = bit_depth
456478
457479 def get_components (self ) -> VideoComponents :
458480 return VideoComponents (
@@ -461,18 +483,26 @@ def get_components(self) -> VideoComponents:
461483 frame_rate = self .__components .frame_rate ,
462484 )
463485
486+ def get_bit_depth (self ) -> int :
487+ return self .__bit_depth
488+
464489 def save_to (
465490 self ,
466491 path : str ,
467492 format : VideoContainer = VideoContainer .AUTO ,
468493 codec : VideoCodec = VideoCodec .AUTO ,
469494 metadata : Optional [dict ] = None ,
495+ bit_depth : int | None = None ,
470496 ):
471497 """Save the video to a file path or BytesIO buffer."""
472498 if format != VideoContainer .AUTO and format != VideoContainer .MP4 :
473499 raise ValueError ("Only MP4 format is supported for now" )
474500 if codec != VideoCodec .AUTO and codec != VideoCodec .H264 :
475501 raise ValueError ("Only H264 codec is supported for now" )
502+ # None means "use the depth this video was created with" (CreateVideo's choice).
503+ if bit_depth is None :
504+ bit_depth = self .__bit_depth
505+ is_10bit = bit_depth >= 10
476506 extra_kwargs = {}
477507 if isinstance (format , VideoContainer ) and format != VideoContainer .AUTO :
478508 extra_kwargs ["format" ] = format .value
@@ -488,10 +518,11 @@ def save_to(
488518
489519 frame_rate = Fraction (round (self .__components .frame_rate * 1000 ), 1000 )
490520 # Create a video stream
521+ pix_fmt = "yuv420p10le" if is_10bit else "yuv420p"
491522 video_stream = output .add_stream ('h264' , rate = frame_rate )
492523 video_stream .width = self .__components .images .shape [2 ]
493524 video_stream .height = self .__components .images .shape [1 ]
494- video_stream .pix_fmt = 'yuv420p'
525+ video_stream .pix_fmt = pix_fmt
495526
496527 # Create an audio stream
497528 audio_sample_rate = 1
@@ -505,9 +536,14 @@ def save_to(
505536
506537 # Encode video
507538 for i , frame in enumerate (self .__components .images ):
508- img = (frame * 255 ).clamp (0 , 255 ).byte ().cpu ().numpy () # shape: (H, W, 3)
509- frame = av .VideoFrame .from_ndarray (img , format = 'rgb24' )
510- frame = frame .reformat (format = 'yuv420p' ) # Convert to YUV420P as required by h264
539+ if is_10bit :
540+ # 16-bit RGB keeps float precision through the conversion to 10-bit YUV.
541+ img = (frame .float () * 65535 ).clamp (0 , 65535 ).cpu ().numpy ().astype (np .uint16 ) # shape: (H, W, 3)
542+ frame = av .VideoFrame .from_ndarray (img , format = "rgb48le" )
543+ else :
544+ img = (frame * 255 ).clamp (0 , 255 ).byte ().cpu ().numpy () # shape: (H, W, 3)
545+ frame = av .VideoFrame .from_ndarray (img , format = 'rgb24' )
546+ frame = frame .reformat (format = pix_fmt )
511547 packet = video_stream .encode (frame )
512548 output .mux (packet )
513549
0 commit comments