606 lines
28 KiB
Python
606 lines
28 KiB
Python
import pygame
|
||
import pygame.camera
|
||
import pygame_gui
|
||
from pygame_gui.core import ObjectID
|
||
from torchvision import models
|
||
import torchvision.transforms as T
|
||
import torch
|
||
import numpy as np
|
||
import json
|
||
import pprint
|
||
import torch.nn.functional as F
|
||
|
||
torch.cuda.set_device(0)
|
||
torch.backends.cudnn.benchmark = True
|
||
|
||
torch.backends.cudnn.allow_tf32 = True
|
||
torch.backends.cuda.matmul.allow_tf32 = True
|
||
torch.autograd.set_detect_anomaly(True)
|
||
torch.autograd.profiler.profile (enabled = False)
|
||
torch.backends.cudnn.deterministic = True
|
||
torch.backends.cudnn.benchmark = False
|
||
|
||
# Загрузка предварительно обученной модели DeepLabV3
|
||
device = torch.device('cuda')
|
||
print(device)
|
||
model = models.segmentation.fcn_resnet101(pretrained=True, aux_loss=True).to(device)
|
||
model.eval()
|
||
|
||
# bg_image = pygame.image.load('BG.png')
|
||
|
||
|
||
for model_name in dir(models.segmentation):
|
||
if not model_name.startswith('__'):
|
||
print(model_name)
|
||
|
||
|
||
APP_TITLE = "PyGPU-VBG-0.1"
|
||
THEME_FILE = "theme.json"
|
||
OPENING_CAMERA = "Opening camera"
|
||
CLOSING_CAMERA = "Stopping camera"
|
||
NO_CAMERAS_FOUND = "No cameras found"
|
||
PICK_WEBCAM = "Pick a camera"
|
||
CLOSE_APP = "Close app"
|
||
MASK_SCALE_FACTOR = "Mask scaling ratio:"
|
||
MASK_COLOR_LABEL = "Mask color:"
|
||
CAM_COLOR_BALANCE_LABEL = "Webcam color balance:"
|
||
BRIGHTNESS_LABEL = "Brightness:"
|
||
CONTRAST_LABEL = "Contrast:"
|
||
FAILED_TO_STOP_CAMERA = "Couldn't stop camera"
|
||
NO_WEBCAM_SELECTED = 'No webcam selected'
|
||
|
||
WHITE = (255, 255, 255)
|
||
GRAY = (205, 205, 205)
|
||
BLACK = (0, 0, 0)
|
||
BACKGROUND = (21, 25, 30)
|
||
|
||
RESET_EMOJI = "↺"
|
||
|
||
FONT_NAME = "Unifont"
|
||
FONT_PATH = "./fonts/unifont-15.1.05.otf"
|
||
|
||
SETTINGS_FILE = "settings.json"
|
||
|
||
# Инициализация Pygame и камеры
|
||
pygame.init()
|
||
pygame.camera.init()
|
||
|
||
# Номер камеры и размеры изображения
|
||
|
||
camera = None
|
||
cam_resolution = (640, 480)
|
||
window_size = (870, 480)
|
||
x_offset = -15
|
||
|
||
default_mask_color = [0, 255, 0]
|
||
default_mask_scale = 1
|
||
default_mask_alpha = 255
|
||
default_cam_contrast = 1
|
||
default_cam_brightness = 0
|
||
default_white_balance_r = 1
|
||
default_white_balance_g = 1
|
||
default_white_balance_b = 1
|
||
|
||
mask_color = default_mask_color
|
||
mask_scale = default_mask_scale
|
||
mask_alpha = default_mask_alpha
|
||
cam_contrast = default_cam_contrast
|
||
cam_brightness = default_cam_brightness
|
||
white_balance_r = default_white_balance_r
|
||
white_balance_g = default_white_balance_g
|
||
white_balance_b = default_white_balance_b
|
||
current_camera = 0
|
||
#################################
|
||
alpha = 0.05
|
||
smoothed_mask = None
|
||
#################################
|
||
|
||
# Создание объекта камеры
|
||
cam_list = pygame.camera.list_cameras()
|
||
if not cam_list:
|
||
raise ValueError(NO_CAMERAS_FOUND)
|
||
|
||
# Экран Pygame
|
||
screen = pygame.display.set_mode(window_size, pygame.RESIZABLE)
|
||
pygame.display.set_caption(APP_TITLE)
|
||
# pygame.display.set_minimum_dimensions((870, 480))
|
||
ui_manager = pygame_gui.UIManager(window_size, THEME_FILE)
|
||
ui_manager.add_font_paths(FONT_NAME, FONT_PATH)
|
||
ui_manager.print_unused_fonts()
|
||
|
||
|
||
ui_camera_list = cam_list.copy()
|
||
ui_camera_list.insert(0, PICK_WEBCAM)
|
||
|
||
y_offset = 0
|
||
default_offset_increment = 25
|
||
|
||
class CustomRect(pygame.Rect):
|
||
def __init__(self, x, y, width, height, bottomright):
|
||
super().__init__(x, y, width, height)
|
||
self.bottomright = bottomright
|
||
|
||
class align_top_right(pygame.Rect):
|
||
def __init__(self, offset, size):
|
||
super().__init__(0, 0, size[0], size[1])
|
||
self.topright = (offset[0], offset[1])
|
||
|
||
|
||
def get_y_offset(*args):
|
||
global y_offset
|
||
height = None
|
||
if len(args) > 0:
|
||
height = args[0]
|
||
|
||
if height is not None:
|
||
y_offset += height
|
||
else:
|
||
y_offset += default_offset_increment
|
||
|
||
return y_offset
|
||
|
||
def camera_name_to_index(name):
|
||
index = 0
|
||
print("Searching for a cam:", name)
|
||
for cam in ui_camera_list:
|
||
print(index, cam)
|
||
if cam == name:
|
||
print("Found cam at index:",index)
|
||
return index
|
||
index += 1
|
||
|
||
|
||
|
||
|
||
no_camera_selected_label = pygame_gui.elements.UILabel(
|
||
text = NO_WEBCAM_SELECTED,
|
||
relative_rect=pygame.Rect((cam_resolution[1]-180)/2, window_size[1]/2-12, cam_resolution[1]-150, 24),
|
||
object_id=ObjectID(class_id='@custom_labels', object_id='#no_webcam_label'),
|
||
manager=ui_manager)
|
||
|
||
slider_threshold = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, 0), (180, 25)),
|
||
manager=ui_manager,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
start_value=1,
|
||
value_range=(1, 20))
|
||
|
||
|
||
camera_select = pygame_gui.elements.UIDropDownMenu(
|
||
options_list = ui_camera_list,
|
||
starting_option = PICK_WEBCAM,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
relative_rect=align_top_right((x_offset, get_y_offset()), (210, 25)), manager=ui_manager)
|
||
|
||
|
||
slider_label = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset, get_y_offset()), (210, 25)),
|
||
text=MASK_SCALE_FACTOR,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
|
||
mask_scale_slider = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset()), (180, 25)),
|
||
manager=ui_manager,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
start_value=mask_scale*100,
|
||
value_range=(1, 200))
|
||
|
||
slider_label3 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset, get_y_offset()), (210, 25)),
|
||
text=MASK_COLOR_LABEL,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
|
||
slider_label3 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset-180, get_y_offset()), (25, 25)),
|
||
text='r',
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_r = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset(0)), (155, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_mask_color[0],
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(0, 255))
|
||
|
||
slider_label3 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset-180, get_y_offset()), (25, 25)),
|
||
text='g',
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_g = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset(0)), (155, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_mask_color[1],
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(0, 255))
|
||
|
||
slider_label3 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset-180, get_y_offset()), (25, 25)),
|
||
text='b',
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_b = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset(0)), (155, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_mask_color[2],
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(0, 255))
|
||
|
||
slider_label3 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset-180, get_y_offset()), (25, 25)),
|
||
text='a',
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_a = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset(0)), (155, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_mask_alpha,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(0, 255))
|
||
|
||
|
||
##############
|
||
slider_label4 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset-30, get_y_offset()), (180, 25)),
|
||
text=CAM_COLOR_BALANCE_LABEL,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
|
||
slider_label5 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset-180, get_y_offset()), (25, 25)),
|
||
text='r',
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_wbr = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset(0)), (155, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_white_balance_r*100,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(0, 200))
|
||
|
||
slider_label6 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset-180, get_y_offset()), (25, 25)),
|
||
text='g',
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_wbg = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset(0)), (155, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_white_balance_g*100,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(0, 200))
|
||
|
||
slider_label7 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset-180, get_y_offset()), (25, 25)),
|
||
text='b',
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_wbb = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset(0)), (155, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_white_balance_b*100,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(0, 200))
|
||
|
||
|
||
slider_label6 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset, get_y_offset()), (210, 25)),
|
||
text=BRIGHTNESS_LABEL,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_brightness = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset()), (180, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_cam_brightness,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(-100, 100))
|
||
|
||
slider_label7 = pygame_gui.elements.UILabel(relative_rect=align_top_right((x_offset, get_y_offset()), (210, 25)),
|
||
text=CONTRAST_LABEL,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
manager=ui_manager)
|
||
slider_contrast = pygame_gui.elements.UIHorizontalSlider(relative_rect=align_top_right((x_offset-30, get_y_offset()), (180, 25)),
|
||
manager=ui_manager,
|
||
start_value=default_cam_contrast*100,
|
||
anchors={'right': 'right', 'top': 'top'},
|
||
value_range=(1, 200))
|
||
|
||
button = pygame_gui.elements.UIButton(relative_rect = CustomRect( 0, 0, 210, 25, (-15, -10) ),
|
||
text=CLOSE_APP,
|
||
anchors={'right': 'right', 'bottom': 'bottom'},
|
||
manager=ui_manager)
|
||
|
||
y_offset = 50
|
||
rbutton1 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbutton2 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbutton3 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbutton4 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbutton5 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbutton6 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbutton7 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbutton8 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbutton9 = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbuttona = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset()), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbuttonb = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset(50)), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
rbuttonc = pygame_gui.elements.UIButton(relative_rect=align_top_right((x_offset, get_y_offset(50)), (25, 25)), text=RESET_EMOJI, manager=ui_manager, anchors={'right': 'right', 'top': 'top'})
|
||
|
||
#######################################
|
||
# Save the JSON object to a file
|
||
def save_data(data, filename):
|
||
with open(filename, 'w') as f:
|
||
json.dump(data, f, indent=4)
|
||
|
||
# Load the JSON object from a file
|
||
def load_data(filename):
|
||
try:
|
||
with open(filename, 'r') as f:
|
||
return json.load(f)
|
||
except FileNotFoundError:
|
||
print("File not found.")
|
||
return None
|
||
|
||
current_settings = {
|
||
|
||
}
|
||
|
||
loaded_data = load_data(SETTINGS_FILE)
|
||
|
||
if loaded_data:
|
||
print(loaded_data)
|
||
# cam_resolution = loaded_data['cam_resolution']
|
||
# window_size = loaded_data['window_size']
|
||
alpha = loaded_data['alpha']
|
||
slider_brightness.set_current_value(loaded_data['cam_brightness'])
|
||
slider_contrast.set_current_value(loaded_data['cam_contrast'])
|
||
slider_r.set_current_value(loaded_data['mask_color'][0])
|
||
slider_g.set_current_value(loaded_data['mask_color'][1])
|
||
slider_b.set_current_value(loaded_data['mask_color'][2])
|
||
slider_a.set_current_value(loaded_data['mask_alpha'])
|
||
slider_wbr.set_current_value(loaded_data['white_balance_r'])
|
||
slider_wbg.set_current_value(loaded_data['white_balance_g'])
|
||
slider_wbb.set_current_value(loaded_data['white_balance_b'])
|
||
mask_scale_slider.set_current_value(loaded_data['mask_scale'])
|
||
|
||
|
||
#######################################
|
||
|
||
def get_scaled_image_rect(image_size, dest_rect_size):
|
||
image_width, image_height = image_size
|
||
dest_width, dest_height = dest_rect_size
|
||
|
||
image_aspect_ratio = image_width / image_height
|
||
dest_aspect_ratio = dest_width / dest_height
|
||
|
||
if image_aspect_ratio > dest_aspect_ratio:
|
||
scaled_width = dest_width
|
||
scaled_height = int(dest_width / image_aspect_ratio)
|
||
else:
|
||
scaled_width = int(dest_height * image_aspect_ratio)
|
||
scaled_height = dest_height
|
||
|
||
x = (dest_width - scaled_width) // 2
|
||
y = (dest_height - scaled_height) // 2
|
||
|
||
return pygame.Rect(x, y, scaled_width, scaled_height)
|
||
|
||
|
||
|
||
|
||
|
||
|
||
# Функция предобработки изображения
|
||
def preprocess(image, target_size):
|
||
transform = T.Compose([
|
||
T.ToTensor(),
|
||
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||
T.Resize(target_size)
|
||
])
|
||
image = pygame.surfarray.array3d(image)
|
||
image2 = image.transpose([0, 1, 2]) # Correct the orientation
|
||
image2 = transform(image).unsqueeze(0).to(device)
|
||
return image2
|
||
|
||
def process_image(image):
|
||
|
||
# получение значения настроек изображения с камеры
|
||
white_balance_r = slider_wbr.get_current_value() / 100
|
||
white_balance_g = slider_wbg.get_current_value() / 100
|
||
white_balance_b = slider_wbb.get_current_value() / 100
|
||
|
||
cam_brightness = slider_brightness.get_current_value()
|
||
cam_contrast = slider_contrast.get_current_value() / 100
|
||
|
||
|
||
array = pygame.surfarray.array3d(image)
|
||
|
||
# Изменение яркости
|
||
if cam_brightness > 0 :
|
||
array += cam_brightness
|
||
else :
|
||
array -= cam_brightness*-1
|
||
array = np.clip(array, 0, 255)
|
||
|
||
# Изменение контраста
|
||
array = array * cam_contrast
|
||
array = np.clip(array, 0, 255)
|
||
|
||
|
||
# Изменение баланса белого
|
||
array[:, :, 0] = array[:, :, 0] * white_balance_r
|
||
array[:, :, 1] = array[:, :, 1] * white_balance_g
|
||
array[:, :, 2] = array[:, :, 2] * white_balance_b
|
||
array = np.clip(array, 0, 255)
|
||
# Преобразование массива обратно в изображение
|
||
output_image = pygame.surfarray.make_surface(array)
|
||
return output_image
|
||
|
||
|
||
# Initialize the temporal smoothing buffer
|
||
temporal_smoothing_buffer = None
|
||
temporal_smoothing_alpha = 0.9 # Adjust this value to control the smoothing strength
|
||
|
||
|
||
def process():
|
||
global camera
|
||
# Захват изображения с камеры
|
||
image = camera.get_image()
|
||
|
||
# получение значения размера маски
|
||
mask_scale = mask_scale_slider.get_current_value() / 100
|
||
|
||
# получение цвета маски
|
||
mask_color = [slider_r.get_current_value(), slider_g.get_current_value(), slider_b.get_current_value()]
|
||
mask_alpha = slider_a.get_current_value()
|
||
|
||
# Размеры маски
|
||
mask_resolution = (int(cam_resolution[0] * mask_scale), int(cam_resolution[1] * mask_scale))
|
||
|
||
# Предобработка изображения и применение модели
|
||
output_image = process_image(image)
|
||
input_tensor = preprocess(output_image, mask_resolution)
|
||
|
||
with torch.no_grad():
|
||
output = model(input_tensor)['out'][0]
|
||
# output = torch.nn.functional.threshold(output, threshold=0.8, value=0)
|
||
mask = output.argmax(0).byte().cpu().detach().numpy()
|
||
mask = np.tile(mask[..., None], [1, 1, 3]) # Convert to 3-channel mask
|
||
|
||
# Apply temporal smoothing
|
||
# global temporal_smoothing_buffer
|
||
# if temporal_smoothing_buffer is None or temporal_smoothing_buffer.shape != mask.shape:
|
||
# temporal_smoothing_buffer = mask
|
||
# else:
|
||
# temporal_smoothing_buffer = (1 - temporal_smoothing_alpha) * temporal_smoothing_buffer + temporal_smoothing_alpha * mask
|
||
|
||
# mask = temporal_smoothing_buffer
|
||
|
||
|
||
# Изменение размера маски, чтобы соответствовать исходному изображению
|
||
threshold = slider_threshold.get_current_value()
|
||
|
||
if True:
|
||
#########################
|
||
mask = output.argmax(0).byte()
|
||
mask = mask.unsqueeze(0).unsqueeze(0) # Add batch and channel dimensions
|
||
dilated_mask = F.max_pool2d(mask.float(), kernel_size=threshold, stride=1, padding=int(threshold/2))
|
||
dilated_mask = (dilated_mask > 0).float()
|
||
dilated_mask = dilated_mask.squeeze(0).squeeze(0) # Remove batch and channel dimensions
|
||
dilated_mask = dilated_mask.cpu().numpy()
|
||
dilated_mask = np.tile(dilated_mask[..., None], [1, 1, 3]) # Repeat for each channel
|
||
|
||
mask = np.where(dilated_mask == 1, mask_color, [0, 0, 0])
|
||
mask = pygame.surfarray.make_surface(mask)
|
||
mask = pygame.transform.scale(mask, cam_resolution)
|
||
mask.set_colorkey((0, 0, 0)) # Set black color as transparent
|
||
mask.set_alpha(mask_alpha) # Set transparency for the mask overlay
|
||
|
||
return output_image, mask
|
||
#########################
|
||
else:
|
||
|
||
mask = np.where(mask != 15, mask_color, [0, 0, 0]) # Apply green color to mask
|
||
mask = pygame.surfarray.make_surface(mask)
|
||
mask = pygame.transform.scale(mask, cam_resolution)
|
||
mask.set_colorkey((0, 0, 0)) # Set black color as transparent
|
||
mask.set_alpha(mask_alpha) # Set transparency for the mask overlay
|
||
|
||
return output_image, mask
|
||
|
||
# Главный цикл
|
||
running = True
|
||
while running:
|
||
current_settings['cam_brightness'] = slider_brightness.get_current_value()
|
||
current_settings['cam_contrast'] = slider_contrast.get_current_value()
|
||
current_settings['cam_resolution'] = cam_resolution
|
||
current_settings['window_size'] = window_size
|
||
current_settings['mask_color'] = [slider_r.get_current_value(), slider_g.get_current_value(), slider_b.get_current_value()]
|
||
current_settings['mask_alpha'] = slider_a.get_current_value()
|
||
current_settings['mask_scale'] = mask_scale_slider.get_current_value()
|
||
current_settings['white_balance_r'] = slider_wbr.get_current_value()
|
||
current_settings['white_balance_g'] = slider_wbg.get_current_value()
|
||
current_settings['white_balance_b'] = slider_wbb.get_current_value()
|
||
current_settings['alpha'] = alpha
|
||
|
||
for event in pygame.event.get():
|
||
if event.type == pygame.VIDEORESIZE:
|
||
window_size = event.size
|
||
no_camera_selected_label.set_position((11, event.h/2-12))
|
||
no_camera_selected_label.set_dimensions((window_size[0]-260, 24))
|
||
ui_manager.set_window_resolution((event.w, event.h))
|
||
|
||
ui_manager.process_events(event)
|
||
if event.type == pygame.QUIT:
|
||
running = False
|
||
elif event.type == pygame.KEYDOWN: # Проверка на событие нажатия клавиши
|
||
if event.key == pygame.K_q: # Проверка, нажата ли клавиша 'q'
|
||
running = False
|
||
elif event.type == pygame_gui.UI_BUTTON_PRESSED: # Проверяем, была ли нажата кнопка
|
||
if event.ui_element == button:
|
||
running = False
|
||
elif event.ui_element == rbutton1:
|
||
mask_scale_slider.set_current_value(default_mask_scale * 100)
|
||
elif event.ui_element == rbutton2:
|
||
slider_r.set_current_value(default_mask_color[0])
|
||
slider_g.set_current_value(default_mask_color[1])
|
||
slider_b.set_current_value(default_mask_color[2])
|
||
slider_a.set_current_value(default_mask_alpha)
|
||
elif event.ui_element == rbutton3:
|
||
slider_r.set_current_value(default_mask_color[0])
|
||
elif event.ui_element == rbutton4:
|
||
slider_g.set_current_value(default_mask_color[1])
|
||
elif event.ui_element == rbutton5:
|
||
slider_b.set_current_value(default_mask_color[2])
|
||
elif event.ui_element == rbutton6:
|
||
slider_a.set_current_value(default_mask_alpha)
|
||
elif event.ui_element == rbutton7:
|
||
slider_wbr.set_current_value(default_white_balance_r * 100)
|
||
slider_wbg.set_current_value(default_white_balance_g * 100)
|
||
slider_wbb.set_current_value(default_white_balance_b * 100)
|
||
elif event.ui_element == rbutton8:
|
||
slider_wbr.set_current_value(default_white_balance_r * 100)
|
||
elif event.ui_element == rbutton9:
|
||
slider_wbg.set_current_value(default_white_balance_g * 100)
|
||
elif event.ui_element == rbuttona:
|
||
slider_wbb.set_current_value(default_white_balance_b * 100)
|
||
elif event.ui_element == rbuttonb:
|
||
slider_brightness.set_current_value(default_cam_brightness)
|
||
elif event.ui_element == rbuttonc:
|
||
slider_contrast.set_current_value(default_cam_contrast*100)
|
||
elif event.type == pygame_gui.UI_DROP_DOWN_MENU_CHANGED:
|
||
if event.ui_element == camera_select:
|
||
if current_camera > 0:
|
||
try:
|
||
print(CLOSING_CAMERA, cam_list[current_camera])
|
||
current_camera = 0
|
||
camera.stop()
|
||
no_camera_selected_label.show()
|
||
|
||
except Exception as e:
|
||
print(FAILED_TO_STOP_CAMERA, event.text)
|
||
|
||
selected_index = camera_name_to_index(event.text)
|
||
if selected_index > 0:
|
||
camera = pygame.camera.Camera(cam_list[selected_index -1], cam_resolution)
|
||
camera.start()
|
||
current_camera = selected_index
|
||
no_camera_selected_label.hide()
|
||
|
||
if current_camera > 0:
|
||
screen.fill(BACKGROUND, (cam_resolution[0]+1, 0, window_size[0]-cam_resolution[0], window_size[1]))
|
||
oimage, mask = process()
|
||
|
||
current_resolution = (window_size[0]-230, window_size[1])
|
||
scaled_rect = get_scaled_image_rect(oimage.get_size(), current_resolution)
|
||
scaled_image = pygame.transform.scale(oimage, scaled_rect.size)
|
||
scaled_mask = pygame.transform.scale(mask, scaled_rect.size)
|
||
|
||
# Вывод изображения на экран
|
||
screen.blit(scaled_image, scaled_rect, special_flags=pygame.BLEND_ALPHA_SDL2)
|
||
# screen.blit(bg_image, scaled_rect, special_flags=pygame.BLEND_SUB)
|
||
# Отображение маски поверх изображения
|
||
screen.blit(scaled_mask, scaled_rect, special_flags=pygame.BLEND_ALPHA_SDL2)
|
||
|
||
|
||
|
||
else:
|
||
screen.fill(BACKGROUND, (0, 0, window_size[0], window_size[1]))
|
||
current_resolution = (window_size[0]-230, window_size[1])
|
||
pygame.draw.rect(screen, GRAY, (10, 10, current_resolution[0]-20, current_resolution[1]-20), 1)
|
||
pygame.draw.line(screen, GRAY, (10, 10), (current_resolution[0]-11, current_resolution[1]-11), 1)
|
||
pygame.draw.line(screen, GRAY, (current_resolution[0]-11, 10), (10, current_resolution[1]-11), 1)
|
||
# screen.fill(BACKGROUND, (20, cam_resolution[1]/2-12, cam_resolution[0]-40, 24))
|
||
|
||
ui_manager.update(1 / 60)
|
||
ui_manager.draw_ui(screen)
|
||
pygame.display.update()
|
||
pygame.time.Clock().tick(60)
|
||
pygame.display.flip()
|
||
|
||
# Сохраненеи натроек при закрытии
|
||
save_data(current_settings, SETTINGS_FILE)
|
||
|
||
# Освобождение ресурсов
|
||
camera.stop()
|
||
pygame.quit()
|
||
|