File size: 4,783 Bytes
b171568 |
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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
import json
from pathlib import Path
import cv2
def get_video_info(video_path, prompt_text):
"""Extract video information using OpenCV and corresponding prompt text"""
cap = cv2.VideoCapture(str(video_path))
if not cap.isOpened():
print(f"Error: Could not open video {video_path}")
return None
# Get video properties
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
duration = frame_count / fps if fps > 0 else 0
cap.release()
return {
"path": video_path.name,
"resolution": {
"width": width,
"height": height
},
"fps": fps,
"duration": duration,
"cap": [prompt_text]
}
def read_prompt_file(prompt_path):
"""Read and return the content of a prompt file"""
try:
with open(prompt_path, 'r', encoding='utf-8') as f:
return f.read().strip()
except Exception as e:
print(f"Error reading prompt file {prompt_path}: {e}")
return None
def process_videos_and_prompts(video_dir_path, prompt_dir_path, verbose=False):
"""Process videos and their corresponding prompt files
Args:
video_dir_path (str): Path to directory containing video files
prompt_dir_path (str): Path to directory containing prompt files
verbose (bool): Whether to print verbose processing information
"""
video_dir = Path(video_dir_path)
prompt_dir = Path(prompt_dir_path)
processed_data = []
# Ensure directories exist
if not video_dir.exists() or not prompt_dir.exists():
print(
f"Error: One or both directories do not exist:\nVideos: {video_dir}\nPrompts: {prompt_dir}"
)
return []
# Process each video file
for video_file in video_dir.glob('*.mp4'):
video_name = video_file.stem
prompt_file = prompt_dir / f"{video_name}.txt"
# Check if corresponding prompt file exists
if not prompt_file.exists():
print(f"Warning: No prompt file found for video {video_name}")
continue
# Read prompt content
prompt_text = read_prompt_file(prompt_file)
if prompt_text is None:
continue
# Process video and add to results
video_info = get_video_info(video_file, prompt_text)
if video_info:
processed_data.append(video_info)
return processed_data
def save_results(processed_data, output_path):
"""Save processed data to JSON file
Args:
processed_data (list): List of processed video information
output_path (str): Full path for output JSON file
"""
output_path = Path(output_path)
# Create parent directories if they don't exist
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, 'w', encoding='utf-8') as f:
json.dump(processed_data, f, indent=2, ensure_ascii=False)
return output_path
def parse_args():
"""Parse command line arguments"""
import argparse
parser = argparse.ArgumentParser(
description='Process videos and their corresponding prompt files')
parser.add_argument('--video_dir',
'-v',
required=True,
help='Directory containing video files')
parser.add_argument('--prompt_dir',
'-p',
required=True,
help='Directory containing prompt text files')
parser.add_argument(
'--output_path',
'-o',
required=True,
help=
'Full path for output JSON file (e.g., /path/to/output/videos2caption.json)'
)
parser.add_argument('--verbose',
action='store_true',
help='Print verbose processing information')
return parser.parse_args()
if __name__ == "__main__":
# Parse command line arguments
args = parse_args()
# Process videos and prompts
processed_videos = process_videos_and_prompts(args.video_dir,
args.prompt_dir,
args.verbose)
if processed_videos:
# Save results
output_path = save_results(processed_videos, args.output_path)
print(f"\nProcessed {len(processed_videos)} videos")
print(f"Results saved to: {output_path}")
# Print example of processed data
print("\nExample of processed video info:")
print(json.dumps(processed_videos[0], indent=2))
else:
print("No videos were processed successfully")
|