144 lines
5.3 KiB
Python
144 lines
5.3 KiB
Python
import sys
|
|
import cv2
|
|
import numpy as np
|
|
from PyQt6.QtWidgets import QApplication, QLabel, QPushButton, QVBoxLayout, QWidget, QFileDialog, QSlider, QHBoxLayout
|
|
from PyQt6.QtCore import Qt, QTimer, QPoint
|
|
from PyQt6.QtGui import QImage, QPixmap, QMouseEvent
|
|
|
|
|
|
class ChromaPreviewer(QWidget):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.setWindowTitle("Live Chroma Previewer")
|
|
self.setGeometry(100, 100, 800, 600)
|
|
|
|
# Video capture
|
|
self.cap = cv2.VideoCapture(0)
|
|
|
|
# Default chroma key values
|
|
self.hue_center = 50 # Default to greenish hue
|
|
self.tolerance = 10 # Default tolerance
|
|
self.lower_bound = np.array([self.hue_center - self.tolerance, 50, 50], dtype=np.uint8)
|
|
self.upper_bound = np.array([self.hue_center + self.tolerance, 255, 255], dtype=np.uint8)
|
|
self.background = None
|
|
self.bg_video = None
|
|
self.bg_video_cap = None
|
|
|
|
# UI Elements
|
|
self.video_label = QLabel(self)
|
|
self.video_label.mousePressEvent = self.pick_color # Allow picking chroma key color
|
|
self.load_bg_button = QPushButton("Load Background Image/Video")
|
|
self.load_bg_button.clicked.connect(self.load_background)
|
|
|
|
self.tolerance_slider = QSlider(Qt.Orientation.Horizontal)
|
|
self.tolerance_slider.setRange(0, 50)
|
|
self.tolerance_slider.setValue(self.tolerance)
|
|
self.tolerance_slider.valueChanged.connect(self.update_tolerance)
|
|
|
|
self.layout = QVBoxLayout()
|
|
self.layout.addWidget(self.video_label)
|
|
self.layout.addWidget(self.load_bg_button)
|
|
|
|
slider_layout = QHBoxLayout()
|
|
slider_layout.addWidget(QLabel("Tolerance"))
|
|
slider_layout.addWidget(self.tolerance_slider)
|
|
self.layout.addLayout(slider_layout)
|
|
|
|
self.setLayout(self.layout)
|
|
|
|
# Timer for updating frames
|
|
self.timer = QTimer()
|
|
self.timer.timeout.connect(self.update_frame)
|
|
self.timer.start(30)
|
|
|
|
def load_background(self):
|
|
file_path, _ = QFileDialog.getOpenFileName(self, "Select Background Image or Video", "",
|
|
"Images/Videos (*.png *.jpg *.jpeg *.mp4 *.avi *.mov)")
|
|
|
|
if file_path:
|
|
if file_path.lower().endswith(('.mp4', '.avi', '.mov')):
|
|
self.bg_video_cap = cv2.VideoCapture(file_path)
|
|
self.bg_video = True
|
|
self.background = None # Clear image background if video is used
|
|
else:
|
|
self.background = cv2.imread(file_path)
|
|
self.background = cv2.cvtColor(self.background, cv2.COLOR_BGR2RGB)
|
|
self.bg_video = False
|
|
self.bg_video_cap = None
|
|
|
|
def update_tolerance(self, value):
|
|
self.tolerance = value
|
|
self.update_chroma_key()
|
|
|
|
def pick_color(self, event: QMouseEvent):
|
|
x, y = event.pos().x(), event.pos().y()
|
|
if hasattr(self, 'current_frame') and self.current_frame is not None:
|
|
h, w, _ = self.current_frame.shape
|
|
if 0 <= x < w and 0 <= y < h:
|
|
pixel_color = self.current_frame[y, x]
|
|
hsv_color = cv2.cvtColor(np.uint8([[pixel_color]]), cv2.COLOR_RGB2HSV)[0][0]
|
|
|
|
self.hue_center = int(hsv_color[0]) # Update center hue
|
|
self.update_chroma_key()
|
|
|
|
def update_chroma_key(self):
|
|
self.lower_bound = np.array([
|
|
np.clip(self.hue_center - self.tolerance, 0, 179),
|
|
50, 50
|
|
], dtype=np.uint8)
|
|
|
|
self.upper_bound = np.array([
|
|
np.clip(self.hue_center + self.tolerance, 0, 179),
|
|
255, 255
|
|
], dtype=np.uint8)
|
|
|
|
def update_frame(self):
|
|
ret, frame = self.cap.read()
|
|
if not ret:
|
|
return
|
|
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
self.current_frame = frame.copy()
|
|
hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
|
|
|
|
mask = cv2.inRange(hsv, self.lower_bound, self.upper_bound)
|
|
mask_inv = cv2.bitwise_not(mask)
|
|
|
|
fg = cv2.bitwise_and(frame, frame, mask=mask_inv)
|
|
|
|
if self.bg_video and self.bg_video_cap:
|
|
ret_bg, bg_frame = self.bg_video_cap.read()
|
|
if not ret_bg:
|
|
self.bg_video_cap.set(cv2.CAP_PROP_POS_FRAMES, 0) # Loop the video
|
|
ret_bg, bg_frame = self.bg_video_cap.read()
|
|
|
|
if ret_bg:
|
|
bg_frame = cv2.cvtColor(bg_frame, cv2.COLOR_BGR2RGB)
|
|
bg_resized = cv2.resize(bg_frame, (frame.shape[1], frame.shape[0]))
|
|
bg = cv2.bitwise_and(bg_resized, bg_resized, mask=mask)
|
|
frame = cv2.add(fg, bg)
|
|
elif self.background is not None:
|
|
bg_resized = cv2.resize(self.background, (frame.shape[1], frame.shape[0]))
|
|
bg = cv2.bitwise_and(bg_resized, bg_resized, mask=mask)
|
|
frame = cv2.add(fg, bg)
|
|
else:
|
|
frame = fg
|
|
|
|
h, w, ch = frame.shape
|
|
bytes_per_line = ch * w
|
|
qimg = QImage(frame.data, w, h, bytes_per_line, QImage.Format.Format_RGB888)
|
|
self.video_label.setPixmap(QPixmap.fromImage(qimg))
|
|
|
|
def closeEvent(self, event):
|
|
self.cap.release()
|
|
if self.bg_video_cap:
|
|
self.bg_video_cap.release()
|
|
event.accept()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app = QApplication(sys.argv)
|
|
window = ChromaPreviewer()
|
|
window.show()
|
|
sys.exit(app.exec())
|