File size: 4,890 Bytes
83e35a7 | 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 | """
Quality and Color Enhancer - Improves image quality while preserving natural colors
"""
import cv2
import numpy as np
from PIL import Image, ImageEnhance
class QualityColorEnhancer:
def __init__(self):
self.enhance_quality = True
self.enhance_colors = True
def enhance_frame(self, frame_path: str, output_path: str = None) -> str:
"""Enhance frame quality and colors"""
if output_path is None:
output_path = frame_path
try:
# Read image
img = cv2.imread(frame_path)
if img is None:
return frame_path
print(f"🎨 Enhancing {frame_path}...")
# 1. Denoise
img = cv2.fastNlMeansDenoisingColored(img, None, 3, 3, 7, 21)
# 2. Improve sharpness
kernel = np.array([[-1,-1,-1],
[-1, 9,-1],
[-1,-1,-1]])
sharpened = cv2.filter2D(img, -1, kernel)
# Blend with original (to avoid over-sharpening)
img = cv2.addWeighted(img, 0.5, sharpened, 0.5, 0)
# 3. Enhance colors using PIL
img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
# Increase color vibrancy
color_enhancer = ImageEnhance.Color(img_pil)
img_pil = color_enhancer.enhance(1.3) # 30% more colorful
# Skip brightness adjustment - images already bright enough
# brightness_enhancer = ImageEnhance.Brightness(img_pil)
# img_pil = brightness_enhancer.enhance(1.0) # No change
# Adjust contrast - reduced
contrast_enhancer = ImageEnhance.Contrast(img_pil)
img_pil = contrast_enhancer.enhance(1.1) # Only 10% more contrast
# Convert back to OpenCV
img = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
# 4. Auto white balance
img = self._auto_white_balance(img)
# 5. Enhance details in dark areas
img = self._enhance_dark_areas(img)
# Save with high quality
cv2.imwrite(output_path, img, [cv2.IMWRITE_JPEG_QUALITY, 100])
return output_path
except Exception as e:
print(f"⚠️ Enhancement failed: {e}")
return frame_path
def _auto_white_balance(self, img):
"""Simple auto white balance with safety checks"""
try:
result = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
avg_a = np.average(result[:, :, 1])
avg_b = np.average(result[:, :, 2])
# More conservative white balance
result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 0] / 255.0) * 0.5)
result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 0] / 255.0) * 0.5)
result = cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
return result
except:
# If white balance fails, return original
return img
def _enhance_dark_areas(self, img):
"""Enhance details in dark areas without affecting bright areas"""
# Create a mask for dark areas
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
mask = cv2.inRange(gray, 0, 100) # Dark areas
# Apply CLAHE only to dark areas
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
l_channel = lab[:, :, 0]
# Create CLAHE object
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
enhanced_l = clahe.apply(l_channel)
# Apply enhancement only to masked areas
l_channel = np.where(mask > 0, enhanced_l, l_channel)
lab[:, :, 0] = l_channel
# Convert back
result = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
return result
def batch_enhance(self, frames_dir: str):
"""Enhance all frames in directory"""
import os
frames = sorted([f for f in os.listdir(frames_dir) if f.endswith('.png')])
print(f"🎨 Enhancing {len(frames)} frames for better quality and colors...")
for i, frame in enumerate(frames):
# Skip last frame if it's problematic
is_last = (i == len(frames) - 1)
if is_last:
print(f" ⚠️ Skipping last frame {frame} to preserve original colors")
continue
frame_path = os.path.join(frames_dir, frame)
self.enhance_frame(frame_path)
print(f" ✓ Enhanced {i+1}/{len(frames)}")
print("✅ Quality and color enhancement complete!") |