vidaio-subnet / services /compress /utils /processing_utils.py

Ahmad Ayad

added passing for target codec for both synthetics and organics (#94)

4e4aa87 11 months ago

12.2 kB

import os
import time
import json
import shutil
import subprocess
from .encode_video import encode_video
from .classify_scene import classify_scene_with_model, extract_frames_from_scene
from .analyze_video_fast import analyze_video_fast

def get_video_bitrate(video_path):
    """Get the bitrate of a video file in kbps."""
    try:
        cmd = [
            'ffprobe', '-v', 'quiet', '-print_format', 'json', 
            '-show_format', video_path
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        data = json.loads(result.stdout)
        
        bitrate = data.get('format', {}).get('bit_rate')
        if bitrate:
            return float(bitrate) / 1000  # Convert to kbps
        return None
    except Exception as e:
        print(f"Error getting bitrate for {video_path}: {e}")
        return None

def analyze_input_compression(video_path):
    """Analyze if input video is already well compressed."""
    try:
        cmd = [
            'ffprobe', '-v', 'quiet', '-print_format', 'json', 
            '-show_streams', '-show_format', video_path
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        data = json.loads(result.stdout)
        
        video_stream = next((s for s in data.get('streams', []) if s.get('codec_type') == 'video'), None)
        if not video_stream:
            return None
            
        codec = video_stream.get('codec_name', '')
        bitrate = data.get('format', {}).get('bit_rate')
        duration = float(data.get('format', {}).get('duration', 0))
        
        # Calculate pixels per second
        width = int(video_stream.get('width', 0))
        height = int(video_stream.get('height', 0))
        fps_str = video_stream.get('r_frame_rate', '0/1')
        fps = eval(fps_str) if '/' in fps_str else float(fps_str)
        
        pixels_per_second = width * height * fps
        
        return {
            'codec': codec,
            'bitrate_kbps': float(bitrate) / 1000 if bitrate else None,
            'duration': duration,
            'resolution': (width, height),
            'fps': fps,
            'pixels_per_second': pixels_per_second,
            'bits_per_pixel': (float(bitrate) / pixels_per_second) if bitrate and pixels_per_second > 0 else None
        }
    except Exception as e:
        print(f"Error analyzing input compression: {e}")
        return None

def should_skip_encoding(input_path, estimated_cq, target_vmaf, logging_enabled=True):
    """Determine if encoding should be skipped to prevent size increase."""
    
    input_analysis = analyze_input_compression(input_path)
    if not input_analysis:
        return False, estimated_cq + 1
    
    input_bitrate = input_analysis.get('bitrate_kbps', 0)
    bits_per_pixel = input_analysis.get('bits_per_pixel', 0)
    codec = input_analysis.get('codec', '').lower()
    
    if bits_per_pixel and bits_per_pixel < 0.02 and target_vmaf < 90:  
        if logging_enabled:
            print(f"⚠️ Input very well compressed (bits/pixel: {bits_per_pixel:.4f}). Encoding anyway with conservative CQ.")
        return False, min(estimated_cq + 8, 51)
    
    if target_vmaf >= 90:
        if logging_enabled:
            print(f"🎯 High target VMAF ({target_vmaf}) - encoding anyway with adjusted CQ")
        return False, estimated_cq
    
    adjusted_cq = calculate_bitrate_aware_cq(estimated_cq, input_analysis, target_vmaf, "unknown")
    
    if input_bitrate > 10000:  
        adjusted_cq = min(adjusted_cq + 4, 51)
        if logging_enabled:
            print(f"Very high bitrate input ({input_bitrate:.0f} kbps) - adding safety margin. CQ: {adjusted_cq}")
    
    return False, adjusted_cq 

def encode_scene_with_size_check(scene_path, output_path, codec, adjusted_cq, content_type, contrast_value, codec_mode=None, target_bitrate=None, max_retries=2, logging_enabled=True):
    
    original_size = os.path.getsize(scene_path)
    input_analysis = analyze_input_compression(scene_path)
    
    # Pre-check if encoding should be skipped
    skip_encoding, final_cq = should_skip_encoding(scene_path, adjusted_cq, 93, logging_enabled)
    
    if skip_encoding:
        if logging_enabled:
            print("Skipping encoding - input already optimally compressed")
        return None, 0.0
    
    if final_cq != adjusted_cq:
        if logging_enabled:
            print(f"Adjusted CQ from {adjusted_cq} to {final_cq} based on input analysis")
        adjusted_cq = final_cq
    
    input_bitrate = input_analysis.get('bitrate_kbps', 0) if input_analysis else 0
    if input_bitrate > 4000:
        max_size_ratio = 1.15  # Allow 15% increase for high-bitrate content
    else:
        max_size_ratio = 1.05  # Only allow 5% for normal content
    
    for attempt in range(max_retries + 1):
        if logging_enabled and attempt > 0:
            print(f"Retry attempt {attempt} - increasing CQ to prevent size increase")
        
        # FIXED: Actually increase CQ on retries
        retry_cq = adjusted_cq + (attempt * 4)  # Start with base CQ, then +4, +8
        retry_cq = min(retry_cq, 51)  # Cap at 51
        
        if logging_enabled and attempt > 0:
            print(f"Using CQ {retry_cq} for attempt {attempt}")
        
        start_time = time.time()
        _, encoding_time = encode_video(
            input_path=scene_path,
            output_path=output_path,
            codec=codec,
            rate=retry_cq,  # FIXED: Use retry_cq instead of adjusted_cq
            scene_type=content_type if attempt == 0 else None,
            contrast_value=contrast_value if attempt == 0 else None,
            codec_mode=codec_mode,
            target_bitrate=target_bitrate,
            logging_enabled=logging_enabled
        )
        
        if not os.path.exists(output_path):
            continue
            
        output_size = os.path.getsize(output_path)
        size_ratio = output_size / original_size
        
        if logging_enabled:
            print(f"Size ratio: {size_ratio:.2f} (output/input) with CQ {retry_cq}")
        
        # Accept the result if within threshold OR it's the final attempt
        if size_ratio <= max_size_ratio or attempt == max_retries:
            if size_ratio > max_size_ratio and logging_enabled:
                print(f"⚠️  Accepting oversized output ({size_ratio:.2f}x) - final attempt")
            return output_path, encoding_time
        else:
            if logging_enabled:
                print(f"Output too large ({size_ratio:.2f}x), retrying with higher CQ")
            os.remove(output_path)  # Remove oversized file
    
    return None, 0.0

def classify_scene_from_path(scene_path, temp_dir, scene_classifier_model, available_metrics, 
                           device='cpu', class_mapping=None, logging_enabled=True, 
                           num_frames=3, metrics_scaler=None):
    """
    Classify a scene directly from video path - wrapper that uses enhanced functions
    
    This function is a simplified wrapper that leverages the robust implementations
    
    Args:
        scene_path: Path to the video file
        temp_dir: Temporary directory for frame extraction
        scene_classifier_model: Trained scene classifier model
        available_metrics: List of available metrics for the model
        device: Device to run inference on ('cpu' or 'cuda')
        class_mapping: Class mapping dictionary (optional)
        logging_enabled: Whether to enable logging
        num_frames: Number of frames to extract for classification
        metrics_scaler: Metrics scaler (optional)
        
    Returns:
        tuple: (classification_label, detailed_results, video_features)
        classification_label: Scene classification result
        detailed_results: Classification confidence and probabilities
        video_features: Video analysis metrics from analyze_video_fast
    """
    try:
        
        
        # Step 1: Get video duration using ffprobe
        if logging_enabled:
            print(f"Extracting {num_frames} frames from {scene_path}")
        
        try:
            import subprocess
            cmd = [
                'ffprobe', '-v', 'quiet', '-show_entries', 'format=duration',
                '-of', 'default=noprint_wrappers=1:nokey=1', scene_path
            ]
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            duration = float(result.stdout.strip())
        except Exception as e:
            if logging_enabled:
                print(f"Warning: Could not get video duration: {e}, using default 30s")
            duration = 30.0  # Fallback duration
        
        # Step 2: Create temp directory for frame extraction
        import tempfile
        frame_temp_dir = os.path.join(temp_dir, f"frames_{int(time.time())}")
        os.makedirs(frame_temp_dir, exist_ok=True)
        
        # Step 3: Extract frames using the robust enhanced function
        frame_paths = extract_frames_from_scene(
            video_path=scene_path,
            start_time=0.0,
            end_time=duration,
            num_frames=num_frames,
            output_dir=frame_temp_dir
        )
        
        if not frame_paths:
            if logging_enabled:
                print("Warning: No frames could be extracted from video")
            return "unclear", {
                'confidence_score': 0.0,
                'prob_screen_content': 0.0,
                'prob_animation': 0.0,
                'prob_faces': 0.0,
                'prob_gaming': 0.0,
                'prob_other': 0.0,
                'prob_unclear': 1.0,
                'frame_predictions': [],
                'error': 'Frame extraction failed'
            }, {}
        
        # Step 4: Extract video features
        if logging_enabled:
            print(f"Extracting video features from {scene_path}")
        
        try:
            
            video_features = analyze_video_fast(scene_path, logging_enabled=False)
            
            if not video_features or 'error' in video_features:
                # Create default video features if extraction fails
                if logging_enabled:
                    print("Warning: Video feature extraction failed, using defaults")
                video_features = {}
                for metric in available_metrics:
                    video_features[metric] = 0.5  # Default neutral values
        except ImportError:
            if logging_enabled:
                print("Warning: analyze_video_fast not available, using default features")
            # Create default video features
            video_features = {}
            for metric in available_metrics:
                video_features[metric] = 0.5  # Default neutral values
        
        # Step 5: Classify using the enhanced classification function
        classification_label, detailed_results = classify_scene_with_model(
            frame_paths=frame_paths,
            video_features=video_features,
            scene_classifier=scene_classifier_model,
            metrics_scaler=metrics_scaler,
            available_metrics=available_metrics,
            device=device,
            logging_enabled=logging_enabled
        )
        
        # Cleanup temporary frames
        try:
            shutil.rmtree(frame_temp_dir)
        except Exception as e:
            if logging_enabled:
                print(f"Warning: Could not cleanup temp frames: {e}")
        
        return classification_label, detailed_results, video_features
        
    except Exception as e:
        if logging_enabled:
            print(f"Error in classify_scene_from_path: {e}")
        
        # Cleanup on error
        try:
            if 'frame_temp_dir' in locals():
                shutil.rmtree(frame_temp_dir)
        except:
            pass
        
        return "unclear", {
            'confidence_score': 0.0,
            'prob_screen_content': 0.0,
            'prob_animation': 0.0,
            'prob_faces': 0.0,
            'prob_gaming': 0.0,
            'prob_other': 0.0,
            'prob_unclear': 1.0,
            'frame_predictions': [],
            'error': str(e)
        }, {}

	import os
	import time
	import json
	import shutil
	import subprocess
	from .encode_video import encode_video
	from .classify_scene import classify_scene_with_model, extract_frames_from_scene
	from .analyze_video_fast import analyze_video_fast

	def get_video_bitrate(video_path):
	"""Get the bitrate of a video file in kbps."""
	try:
	cmd = [
	'ffprobe', '-v', 'quiet', '-print_format', 'json',
	'-show_format', video_path
	]
	result = subprocess.run(cmd, capture_output=True, text=True, check=True)
	data = json.loads(result.stdout)

	bitrate = data.get('format', {}).get('bit_rate')
	if bitrate:
	return float(bitrate) / 1000 # Convert to kbps
	return None
	except Exception as e:
	print(f"Error getting bitrate for {video_path}: {e}")
	return None

	def analyze_input_compression(video_path):
	"""Analyze if input video is already well compressed."""
	try:
	cmd = [
	'ffprobe', '-v', 'quiet', '-print_format', 'json',
	'-show_streams', '-show_format', video_path
	]
	result = subprocess.run(cmd, capture_output=True, text=True, check=True)
	data = json.loads(result.stdout)

	video_stream = next((s for s in data.get('streams', []) if s.get('codec_type') == 'video'), None)
	if not video_stream:
	return None

	codec = video_stream.get('codec_name', '')
	bitrate = data.get('format', {}).get('bit_rate')
	duration = float(data.get('format', {}).get('duration', 0))

	# Calculate pixels per second
	width = int(video_stream.get('width', 0))
	height = int(video_stream.get('height', 0))
	fps_str = video_stream.get('r_frame_rate', '0/1')
	fps = eval(fps_str) if '/' in fps_str else float(fps_str)

	pixels_per_second = width * height * fps

	return {
	'codec': codec,
	'bitrate_kbps': float(bitrate) / 1000 if bitrate else None,
	'duration': duration,
	'resolution': (width, height),
	'fps': fps,
	'pixels_per_second': pixels_per_second,
	'bits_per_pixel': (float(bitrate) / pixels_per_second) if bitrate and pixels_per_second > 0 else None
	}
	except Exception as e:
	print(f"Error analyzing input compression: {e}")
	return None

	def should_skip_encoding(input_path, estimated_cq, target_vmaf, logging_enabled=True):
	"""Determine if encoding should be skipped to prevent size increase."""

	input_analysis = analyze_input_compression(input_path)
	if not input_analysis:
	return False, estimated_cq + 1

	input_bitrate = input_analysis.get('bitrate_kbps', 0)
	bits_per_pixel = input_analysis.get('bits_per_pixel', 0)
	codec = input_analysis.get('codec', '').lower()

	if bits_per_pixel and bits_per_pixel < 0.02 and target_vmaf < 90:
	if logging_enabled:
	print(f"⚠️ Input very well compressed (bits/pixel: {bits_per_pixel:.4f}). Encoding anyway with conservative CQ.")
	return False, min(estimated_cq + 8, 51)

	if target_vmaf >= 90:
	if logging_enabled:
	print(f"🎯 High target VMAF ({target_vmaf}) - encoding anyway with adjusted CQ")
	return False, estimated_cq

	adjusted_cq = calculate_bitrate_aware_cq(estimated_cq, input_analysis, target_vmaf, "unknown")

	if input_bitrate > 10000:
	adjusted_cq = min(adjusted_cq + 4, 51)
	if logging_enabled:
	print(f"Very high bitrate input ({input_bitrate:.0f} kbps) - adding safety margin. CQ: {adjusted_cq}")

	return False, adjusted_cq

	def encode_scene_with_size_check(scene_path, output_path, codec, adjusted_cq, content_type, contrast_value, codec_mode=None, target_bitrate=None, max_retries=2, logging_enabled=True):

	original_size = os.path.getsize(scene_path)
	input_analysis = analyze_input_compression(scene_path)

	# Pre-check if encoding should be skipped
	skip_encoding, final_cq = should_skip_encoding(scene_path, adjusted_cq, 93, logging_enabled)

	if skip_encoding:
	if logging_enabled:
	print("Skipping encoding - input already optimally compressed")
	return None, 0.0

	if final_cq != adjusted_cq:
	if logging_enabled:
	print(f"Adjusted CQ from {adjusted_cq} to {final_cq} based on input analysis")
	adjusted_cq = final_cq

	input_bitrate = input_analysis.get('bitrate_kbps', 0) if input_analysis else 0
	if input_bitrate > 4000:
	max_size_ratio = 1.15 # Allow 15% increase for high-bitrate content
	else:
	max_size_ratio = 1.05 # Only allow 5% for normal content

	for attempt in range(max_retries + 1):
	if logging_enabled and attempt > 0:
	print(f"Retry attempt {attempt} - increasing CQ to prevent size increase")

	# FIXED: Actually increase CQ on retries
	retry_cq = adjusted_cq + (attempt * 4) # Start with base CQ, then +4, +8
	retry_cq = min(retry_cq, 51) # Cap at 51

	if logging_enabled and attempt > 0:
	print(f"Using CQ {retry_cq} for attempt {attempt}")

	start_time = time.time()
	_, encoding_time = encode_video(
	input_path=scene_path,
	output_path=output_path,
	codec=codec,
	rate=retry_cq, # FIXED: Use retry_cq instead of adjusted_cq
	scene_type=content_type if attempt == 0 else None,
	contrast_value=contrast_value if attempt == 0 else None,
	codec_mode=codec_mode,
	target_bitrate=target_bitrate,
	logging_enabled=logging_enabled
	)

	if not os.path.exists(output_path):
	continue

	output_size = os.path.getsize(output_path)
	size_ratio = output_size / original_size

	if logging_enabled:
	print(f"Size ratio: {size_ratio:.2f} (output/input) with CQ {retry_cq}")

	# Accept the result if within threshold OR it's the final attempt
	if size_ratio <= max_size_ratio or attempt == max_retries:
	if size_ratio > max_size_ratio and logging_enabled:
	print(f"⚠️ Accepting oversized output ({size_ratio:.2f}x) - final attempt")
	return output_path, encoding_time
	else:
	if logging_enabled:
	print(f"Output too large ({size_ratio:.2f}x), retrying with higher CQ")
	os.remove(output_path) # Remove oversized file

	return None, 0.0

	def classify_scene_from_path(scene_path, temp_dir, scene_classifier_model, available_metrics,
	device='cpu', class_mapping=None, logging_enabled=True,
	num_frames=3, metrics_scaler=None):
	"""
	Classify a scene directly from video path - wrapper that uses enhanced functions

	This function is a simplified wrapper that leverages the robust implementations

	Args:
	scene_path: Path to the video file
	temp_dir: Temporary directory for frame extraction
	scene_classifier_model: Trained scene classifier model
	available_metrics: List of available metrics for the model
	device: Device to run inference on ('cpu' or 'cuda')
	class_mapping: Class mapping dictionary (optional)
	logging_enabled: Whether to enable logging
	num_frames: Number of frames to extract for classification
	metrics_scaler: Metrics scaler (optional)

	Returns:
	tuple: (classification_label, detailed_results, video_features)
	classification_label: Scene classification result
	detailed_results: Classification confidence and probabilities
	video_features: Video analysis metrics from analyze_video_fast
	"""
	try:


	# Step 1: Get video duration using ffprobe
	if logging_enabled:
	print(f"Extracting {num_frames} frames from {scene_path}")

	try:
	import subprocess
	cmd = [
	'ffprobe', '-v', 'quiet', '-show_entries', 'format=duration',
	'-of', 'default=noprint_wrappers=1:nokey=1', scene_path
	]
	result = subprocess.run(cmd, capture_output=True, text=True, check=True)
	duration = float(result.stdout.strip())
	except Exception as e:
	if logging_enabled:
	print(f"Warning: Could not get video duration: {e}, using default 30s")
	duration = 30.0 # Fallback duration

	# Step 2: Create temp directory for frame extraction
	import tempfile
	frame_temp_dir = os.path.join(temp_dir, f"frames_{int(time.time())}")
	os.makedirs(frame_temp_dir, exist_ok=True)

	# Step 3: Extract frames using the robust enhanced function
	frame_paths = extract_frames_from_scene(
	video_path=scene_path,
	start_time=0.0,
	end_time=duration,
	num_frames=num_frames,
	output_dir=frame_temp_dir
	)

	if not frame_paths:
	if logging_enabled:
	print("Warning: No frames could be extracted from video")
	return "unclear", {
	'confidence_score': 0.0,
	'prob_screen_content': 0.0,
	'prob_animation': 0.0,
	'prob_faces': 0.0,
	'prob_gaming': 0.0,
	'prob_other': 0.0,
	'prob_unclear': 1.0,
	'frame_predictions': [],
	'error': 'Frame extraction failed'
	}, {}

	# Step 4: Extract video features
	if logging_enabled:
	print(f"Extracting video features from {scene_path}")

	try:

	video_features = analyze_video_fast(scene_path, logging_enabled=False)

	if not video_features or 'error' in video_features:
	# Create default video features if extraction fails
	if logging_enabled:
	print("Warning: Video feature extraction failed, using defaults")
	video_features = {}
	for metric in available_metrics:
	video_features[metric] = 0.5 # Default neutral values
	except ImportError:
	if logging_enabled:
	print("Warning: analyze_video_fast not available, using default features")
	# Create default video features
	video_features = {}
	for metric in available_metrics:
	video_features[metric] = 0.5 # Default neutral values

	# Step 5: Classify using the enhanced classification function
	classification_label, detailed_results = classify_scene_with_model(
	frame_paths=frame_paths,
	video_features=video_features,
	scene_classifier=scene_classifier_model,
	metrics_scaler=metrics_scaler,
	available_metrics=available_metrics,
	device=device,
	logging_enabled=logging_enabled
	)

	# Cleanup temporary frames
	try:
	shutil.rmtree(frame_temp_dir)
	except Exception as e:
	if logging_enabled:
	print(f"Warning: Could not cleanup temp frames: {e}")

	return classification_label, detailed_results, video_features

	except Exception as e:
	if logging_enabled:
	print(f"Error in classify_scene_from_path: {e}")

	# Cleanup on error
	try:
	if 'frame_temp_dir' in locals():
	shutil.rmtree(frame_temp_dir)
	except:
	pass

	return "unclear", {
	'confidence_score': 0.0,
	'prob_screen_content': 0.0,
	'prob_animation': 0.0,
	'prob_faces': 0.0,
	'prob_gaming': 0.0,
	'prob_other': 0.0,
	'prob_unclear': 1.0,
	'frame_predictions': [],
	'error': str(e)
	}, {}