Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
iperov
GitHub Repository: iperov/deepfacelab
Path: blob/master/merger/MergerScreen/MergerScreen.py
628 views
1
import math
2
from pathlib import Path
3
4
import numpy as np
5
6
from core import imagelib
7
from core.interact import interact as io
8
from core.cv2ex import *
9
from core import osex
10
11
12
class ScreenAssets(object):
13
waiting_icon_image = cv2_imread ( str(Path(__file__).parent / 'gfx' / 'sand_clock_64.png') )
14
15
@staticmethod
16
def build_checkerboard_a( sh, size=5):
17
h,w = sh[0], sh[1]
18
tile = np.array([[0,1],[1,0]]).repeat(size, axis=0).repeat(size, axis=1)
19
grid = np.tile(tile,(int(math.ceil((h+0.0)/(2*size))),int(math.ceil((w+0.0)/(2*size)))))
20
return grid[:h,:w,None]
21
22
class Screen(object):
23
def __init__(self, initial_scale_to_width=0, initial_scale_to_height=0, image=None, waiting_icon=False, **kwargs):
24
self.initial_scale_to_width = initial_scale_to_width
25
self.initial_scale_to_height = initial_scale_to_height
26
self.image = image
27
self.waiting_icon = waiting_icon
28
29
self.state = -1
30
self.scale = 1
31
self.force_update = True
32
self.is_first_appear = True
33
self.show_checker_board = False
34
35
self.last_screen_shape = (480,640,3)
36
self.checkerboard_image = None
37
self.set_image (image)
38
self.scrn_manager = None
39
40
def set_waiting_icon(self, b):
41
self.waiting_icon = b
42
43
def toggle_show_checker_board(self):
44
self.show_checker_board = not self.show_checker_board
45
self.force_update = True
46
47
def get_image(self):
48
return self.image
49
50
def set_image(self, img):
51
if not img is self.image:
52
self.force_update = True
53
54
self.image = img
55
56
if self.image is not None:
57
self.last_screen_shape = self.image.shape
58
59
if self.initial_scale_to_width != 0:
60
if self.last_screen_shape[1] > self.initial_scale_to_width:
61
self.scale = self.initial_scale_to_width / self.last_screen_shape[1]
62
self.force_update = True
63
self.initial_scale_to_width = 0
64
65
elif self.initial_scale_to_height != 0:
66
if self.last_screen_shape[0] > self.initial_scale_to_height:
67
self.scale = self.initial_scale_to_height / self.last_screen_shape[0]
68
self.force_update = True
69
self.initial_scale_to_height = 0
70
71
72
def diff_scale(self, diff):
73
self.scale = np.clip (self.scale + diff, 0.1, 4.0)
74
self.force_update = True
75
76
def show(self, force=False):
77
new_state = 0 | self.waiting_icon
78
79
if self.state != new_state or self.force_update or force:
80
self.state = new_state
81
self.force_update = False
82
83
if self.image is None:
84
screen = np.zeros ( self.last_screen_shape, dtype=np.uint8 )
85
else:
86
screen = self.image.copy()
87
88
if self.waiting_icon:
89
imagelib.overlay_alpha_image (screen, ScreenAssets.waiting_icon_image, (0,0) )
90
91
h,w,c = screen.shape
92
if self.scale != 1.0:
93
screen = cv2.resize ( screen, ( int(w*self.scale), int(h*self.scale) ) )
94
95
if c == 4:
96
if not self.show_checker_board:
97
screen = screen[...,0:3]
98
else:
99
if self.checkerboard_image is None or self.checkerboard_image.shape[0:2] != screen.shape[0:2]:
100
self.checkerboard_image = ScreenAssets.build_checkerboard_a(screen.shape)
101
102
screen = screen[...,0:3]*0.75 + 64*self.checkerboard_image*(1- (screen[...,3:4].astype(np.float32)/255.0) )
103
screen = screen.astype(np.uint8)
104
105
io.show_image(self.scrn_manager.wnd_name, screen)
106
107
if self.is_first_appear:
108
self.is_first_appear = False
109
#center window
110
desktop_w, desktop_h = osex.get_screen_size()
111
h,w,c = screen.shape
112
cv2.moveWindow(self.scrn_manager.wnd_name, max(0,(desktop_w-w) // 2), max(0, (desktop_h-h) // 2) )
113
114
io.process_messages(0.0001)
115
116
class ScreenManager(object):
117
def __init__(self, window_name="ScreenManager", screens=None, capture_keys=False ):
118
self.screens = screens or []
119
self.current_screen_id = 0
120
121
if self.screens is not None:
122
for screen in self.screens:
123
screen.scrn_manager = self
124
125
self.wnd_name = window_name
126
io.named_window(self.wnd_name)
127
128
129
if capture_keys:
130
io.capture_keys(self.wnd_name)
131
132
def finalize(self):
133
io.destroy_all_windows()
134
135
def get_key_events(self):
136
return io.get_key_events(self.wnd_name)
137
138
def switch_screens(self):
139
self.current_screen_id = (self.current_screen_id + 1) % len(self.screens)
140
self.screens[self.current_screen_id].show(force=True)
141
142
def show_current(self):
143
self.screens[self.current_screen_id].show()
144
145
def get_current(self):
146
return self.screens[self.current_screen_id]
147
148
def set_current(self, screen):
149
self.current_screen_id = self.screens.index(screen)
150
151