File size: 3,466 Bytes
226c7c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from typing import Tuple
import cv2
import magic
import numpy as np
import torch
from cosmos_transfer1.utils import log
# Supported video extensions and corresponding MIME types
SUPPORTED_VIDEO_TYPES = {
".mp4": "video/mp4",
".mkv": "video/x-matroska",
".mov": "video/quicktime",
".avi": "video/x-msvideo",
".webm": "video/webm",
".flv": "video/x-flv",
".wmv": "video/x-ms-wmv",
}
def video_to_tensor(video_path: str, output_path: str, normalize: bool = True) -> Tuple[torch.Tensor, float]:
"""Convert an MP4 video file to a tensor and save it as a .pt file.
Args:
video_path (str): Path to input MP4 video file
output_path (str): Path to save output .pt tensor file
normalize (bool): Whether to normalize pixel values to [-1,1] range (default: True)
Returns:
Tuple[torch.Tensor, float]: Tuple containing:
- Video tensor in shape [C,T,H,W]
- Video FPS
"""
# Open video file
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
raise ValueError(f"Failed to open video file: {video_path}")
# Get video properties
fps = cap.get(cv2.CAP_PROP_FPS)
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# Read first frame to get dimensions
ret, frame = cap.read()
if not ret:
raise ValueError(f"Failed to read frames from video: {video_path}")
height, width = frame.shape[:2]
# Reset video to beginning
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
# Initialize tensor to store frames
frames = []
# Read all frames
while True:
ret, frame = cap.read()
if not ret:
break
# Convert BGR to RGB
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames.append(frame)
cap.release()
log.info(f"frames: {len(frames)}")
# Convert frames to tensor
video_tensor = torch.from_numpy(np.array(frames))
log.info(f"video_tensor shape: {video_tensor.shape}")
# Reshape from [T,H,W,C] to [C,T,H,W]
video_tensor = video_tensor.permute(3, 0, 1, 2)
# Normalize if requested
if normalize:
video_tensor = video_tensor.float() / 127.5 - 1.0
# Save tensor
os.makedirs(os.path.dirname(output_path), exist_ok=True)
torch.save(video_tensor, output_path)
return video_tensor, fps
def is_valid_video(file_path: str) -> bool:
if not os.path.isfile(file_path):
return False
ext = os.path.splitext(file_path)[1].lower()
expected_mime = SUPPORTED_VIDEO_TYPES.get(ext)
if not expected_mime:
return False # Extension not supported
# Detect MIME type from actual file content
detected_mime = magic.from_file(file_path, mime=True)
return detected_mime == expected_mime
|