initial upload
This commit is contained in:
commit
f8fe6c7140
4 changed files with 638 additions and 0 deletions
BIN
fonts/unifont-15.1.05.otf
Normal file
BIN
fonts/unifont-15.1.05.otf
Normal file
Binary file not shown.
9
readme.md
Normal file
9
readme.md
Normal file
|
|
@ -0,0 +1,9 @@
|
||||||
|
Webcam segmentation (fake/virtual green screen) for any OS that supports python3+cuda+torchvision.
|
||||||
|
|
||||||
|
I don't remember all the packages neede to run this, so figure it out on your own (or ask some AI?).
|
||||||
|
|
||||||
|
It eats up to 50% of GPU time on my RTX4070 (regular) and about 500 Mb of vram with image size of 640x480 and produces a stable 30+ fps output.
|
||||||
|
|
||||||
|
Settings are auto-saved.
|
||||||
|
|
||||||
|
WTFPL licence, this is ***NOT*** a good replacement for NVidia broadcast,
|
||||||
23
settings.json
Normal file
23
settings.json
Normal file
|
|
@ -0,0 +1,23 @@
|
||||||
|
{
|
||||||
|
"cam_brightness": 0,
|
||||||
|
"cam_contrast": 100,
|
||||||
|
"cam_resolution": [
|
||||||
|
640,
|
||||||
|
480
|
||||||
|
],
|
||||||
|
"window_size": [
|
||||||
|
870,
|
||||||
|
480
|
||||||
|
],
|
||||||
|
"mask_color": [
|
||||||
|
0,
|
||||||
|
255,
|
||||||
|
0
|
||||||
|
],
|
||||||
|
"mask_alpha": 169,
|
||||||
|
"mask_scale": 48,
|
||||||
|
"white_balance_r": 95,
|
||||||
|
"white_balance_g": 95,
|
||||||
|
"white_balance_b": 94,
|
||||||
|
"alpha": 0.05
|
||||||
|
}
|
||||||
606
test4.py
Normal file
606
test4.py
Normal file
|
|
@ -0,0 +1,606 @@
|
||||||
|
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()
|
||||||
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue