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()