Spaces:
Running
Running
File size: 10,070 Bytes
91fb4ef |
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 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 |
import os
import hashlib
import shutil
from pathlib import Path
import asyncio
import tempfile
import logging
from functools import partial
from typing import Dict, List, Optional, Tuple
import gradio as gr
from scenedetect import detect, ContentDetector, SceneManager, open_video
from scenedetect.video_splitter import split_video_ffmpeg
from config import TRAINING_PATH, STORAGE_PATH, TRAINING_VIDEOS_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH, DEFAULT_PROMPT_PREFIX
from image_preprocessing import detect_black_bars
from video_preprocessing import remove_black_bars
from utils import extract_scene_info, is_video_file, is_image_file, add_prefix_to_caption
logger = logging.getLogger(__name__)
class SplittingService:
def __init__(self):
# Track processing status
self.processing = False
self._current_file: Optional[str] = None
self._scene_counts: Dict[str, int] = {}
self._processing_status: Dict[str, str] = {}
def compute_file_hash(self, file_path: Path) -> str:
"""Compute SHA-256 hash of file"""
sha256_hash = hashlib.sha256()
with open(file_path, "rb") as f:
# Read file in chunks to handle large files
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
return sha256_hash.hexdigest()
def rename_with_hash(self, video_path: Path) -> Tuple[Path, str]:
"""Rename video and caption files using hash
Args:
video_path: Path to video file
Returns:
Tuple of (new video path, hash)
"""
# Compute hash
file_hash = self.compute_file_hash(video_path)
# Rename video file
new_video_path = video_path.parent / f"{file_hash}{video_path.suffix}"
video_path.rename(new_video_path)
# Rename caption file if exists
caption_path = video_path.with_suffix('.txt')
if caption_path.exists():
new_caption_path = caption_path.parent / f"{file_hash}.txt"
caption_path.rename(new_caption_path)
return new_video_path, file_hash
async def process_video(self, video_path: Path, enable_splitting: bool) -> int:
"""Process a single video file to detect and split scenes"""
try:
self._processing_status[video_path.name] = f'Processing video "{video_path.name}"...'
parent_caption_path = video_path.with_suffix('.txt')
# Create output path for split videos
base_name, _ = extract_scene_info(video_path.name)
# Create temporary directory for preprocessed video
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / f"preprocessed_{video_path.name}"
# Try to remove black bars
was_cropped = await asyncio.get_event_loop().run_in_executor(
None,
remove_black_bars,
video_path,
temp_path
)
# Use preprocessed video if cropping was done, otherwise use original
process_path = temp_path if was_cropped else video_path
# Detect scenes if splitting is enabled
if enable_splitting:
video = open_video(str(process_path))
scene_manager = SceneManager()
scene_manager.add_detector(ContentDetector())
scene_manager.detect_scenes(video, show_progress=False)
scenes = scene_manager.get_scene_list()
else:
scenes = []
num_scenes = len(scenes)
if not scenes:
print(f'video "{video_path.name}" is already a single-scene clip')
# captioning is only required if some information is missing
if parent_caption_path.exists():
# if it's a single scene with a caption, we can directly promote it to the training/ dir
#output_video_path = TRAINING_VIDEOS_PATH / f"{base_name}___{1:03d}.mp4"
# WELL ACTUALLY, NOT. The training videos dir removes a lot of thing,
# so it has to stay a "last resort" thing
output_video_path = STAGING_PATH / f"{base_name}___{1:03d}.mp4"
shutil.copy2(process_path, output_video_path)
shutil.copy2(parent_caption_path, output_video_path.with_suffix('.txt'))
parent_caption_path.unlink()
else:
# otherwise it needs to go through the normal captioning process
output_video_path = STAGING_PATH / f"{base_name}___{1:03d}.mp4"
shutil.copy2(process_path, output_video_path)
else:
print(f'video "{video_path.name}" contains {num_scenes} scenes')
# in this scenario, there are multiple subscenes
# even if we have a parent caption, we must caption each of them individually
# the first step is to preserve the parent caption for later use
if parent_caption_path.exists():
output_caption_path = STAGING_PATH / f"{base_name}.txt"
shutil.copy2(parent_caption_path, output_caption_path)
parent_caption_path.unlink()
output_template = str(STAGING_PATH / f"{base_name}___$SCENE_NUMBER.mp4")
# Split video into scenes using the preprocessed video if it exists
await asyncio.get_event_loop().run_in_executor(
None,
lambda: split_video_ffmpeg(
str(process_path),
scenes,
output_file_template=output_template,
show_progress=False
)
)
# Update scene count and status
crop_status = " (black bars removed)" if was_cropped else ""
self._scene_counts[video_path.name] = num_scenes
self._processing_status[video_path.name] = f"{num_scenes} scenes{crop_status}"
# Delete original video
video_path.unlink()
if num_scenes:
gr.Info(f"Extracted {num_scenes} clips from {video_path.name}{crop_status}")
else:
gr.Info(f"Imported {video_path.name}{crop_status}")
return num_scenes
except Exception as e:
self._scene_counts[video_path.name] = 0
self._processing_status[video_path.name] = f"Error: {str(e)}"
raise gr.Error(f"Error processing video {video_path}: {str(e)}")
def get_scene_count(self, video_name: str) -> Optional[int]:
"""Get number of detected scenes for a video
Returns None if video hasn't been scanned
"""
return self._scene_counts.get(video_name)
def get_current_file(self) -> Optional[str]:
"""Get name of file currently being processed"""
return self._current_file
def is_processing(self) -> bool:
"""Check if background processing is running"""
return self.processing
async def start_processing(self, enable_splitting: bool) -> None:
"""Start background processing of unprocessed videos"""
if self.processing:
return
self.processing = True
try:
# Process each video
for video_file in VIDEOS_TO_SPLIT_PATH.glob("*.mp4"):
self._current_file = video_file.name
await self.process_video(video_file, enable_splitting)
finally:
self.processing = False
self._current_file = None
def get_processing_status(self, video_name: str) -> str:
"""Get processing status for a video
Args:
video_name: Name of the video file
Returns:
Status string for the video
"""
if video_name in self._processing_status:
return self._processing_status[video_name]
return "not processed"
def list_unprocessed_videos(self) -> List[List[str]]:
"""List all unprocessed and processed videos with their status.
Images will be ignored.
Returns:
List of lists containing [name, status] for each video
"""
videos = []
# Track processed videos by their base names
processed_videos = {}
for clip_path in STAGING_PATH.glob("*.mp4"):
base_name = clip_path.stem.rsplit('___', 1)[0] + '.mp4'
if base_name in processed_videos:
processed_videos[base_name] += 1
else:
processed_videos[base_name] = 1
# List only video files in processing queue
for video_file in VIDEOS_TO_SPLIT_PATH.glob("*.mp4"):
if is_video_file(video_file): # Only include video files
status = self.get_processing_status(video_file.name)
videos.append([video_file.name, status])
# Add processed videos
for video_name, clip_count in processed_videos.items():
if not (VIDEOS_TO_SPLIT_PATH / video_name).exists():
status = f"Processed ({clip_count} clips)"
videos.append([video_name, status])
return sorted(videos, key=lambda x: (x[1] != "Processing...", x[0].lower()))
|