Skip to content

Commit 50efa62

Browse files
HillobarHillobar
authored
Crystal
1 parent 588fe74 commit 50efa62

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+4533
-0
lines changed

‎rope/Coordinator.py

+248
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
1+
# #!/usr/bin/env python3
2+
import os
3+
import time
4+
5+
import rope.GUI as GUI
6+
import rope.VideoManager as VM
7+
8+
from insightface.app import FaceAnalysis
9+
import onnxruntime
10+
import onnx
11+
12+
import torch
13+
from rope.external.clipseg import CLIPDensePredT
14+
15+
import segmentation_models_pytorch as smp
16+
from collections import OrderedDict
17+
from torchvision import transforms
18+
19+
from rope.external.model import BiSeNet
20+
21+
22+
23+
def coordinator():
24+
global gui, vm, action, frame, r_frame
25+
start = time.time()
26+
27+
if gui.get_action_length() > 0:
28+
action.append(gui.get_action())
29+
if vm.get_action_length() > 0:
30+
action.append(vm.get_action())
31+
##################
32+
if vm.get_frame_length() > 0:
33+
frame.append(vm.get_frame())
34+
35+
if len(frame) > 0:
36+
gui.set_image(frame[0], False)
37+
gui.display_image_in_video_frame()
38+
frame.pop(0)
39+
####################
40+
if vm.get_requested_frame_length() > 0:
41+
r_frame.append(vm.get_requested_frame())
42+
if len(r_frame) > 0:
43+
# print ("1:", time.time())
44+
gui.set_image(r_frame[0], True)
45+
gui.display_image_in_video_frame()
46+
r_frame=[]
47+
####################
48+
if len(action) > 0:
49+
if action[0][0] == "load_target_video":
50+
vm.load_target_video(action[0][1])
51+
#gui.set_slider_position(0)
52+
action.pop(0)
53+
elif action[0][0] == "play_video":
54+
vm.play_video(action[0][1])
55+
action.pop(0)
56+
elif action[0][0] == "set_video_position":
57+
vm.get_requested_video_frame(action[0][1])
58+
action.pop(0)
59+
elif action[0][0] == "find_faces":
60+
gui.find_faces(action[0][1])
61+
action.pop(0)
62+
elif action[0][0] == "clear_faces":
63+
gui.clear_faces()
64+
action.pop(0)
65+
elif action[0][0] == "swap":
66+
if not vm.swapper_model:
67+
gui.set_status("loading Swapper")
68+
swapper, emap = load_swapper_model()
69+
vm.set_swapper_model(swapper, emap)
70+
gui.set_status("Swapper loaded!")
71+
vm.swap_set(action[0][1])
72+
action.pop(0)
73+
elif action[0][0] == "source_embeddings":
74+
vm.load_source_embeddings(action[0][1])
75+
action.pop(0)
76+
elif action[0][0] == "target_faces":
77+
vm.found_faces_assignments = action[0][1]
78+
action.pop(0)
79+
80+
81+
82+
83+
elif action [0][0] == "num_threads":
84+
vm.num_threads = action[0][1]
85+
action.pop(0)
86+
87+
88+
89+
elif action [0][0] == "pos_thresh":
90+
vm.pos_thresh = action[0][1]
91+
action.pop(0)
92+
elif action [0][0] == "neg_thresh":
93+
vm.neg_thresh = action[0][1]
94+
action.pop(0)
95+
elif action [0][0] == "saved_video_path":
96+
vm.saved_video_path = action[0][1]
97+
action.pop(0)
98+
99+
100+
101+
elif action [0][0] == "vid_qual":
102+
vm.vid_qual = int(action[0][1])
103+
action.pop(0)
104+
105+
elif action [0][0] == "parameters":
106+
if action[0][1]["GFPGANState"]:
107+
if not vm.GFPGAN_model:
108+
gui.set_status("loading GFPGAN...")
109+
vm.GFPGAN_model = load_GFPGAN_model()
110+
gui.set_status("GFPGAN loaded!")
111+
if action[0][1]["CLIPState"]:
112+
if not vm.clip_session:
113+
gui.set_status("loading CLIP..")
114+
vm.clip_session, vm.cuda_device = load_clip_model()
115+
gui.set_status("CLIP loaded!")
116+
if action[0][1]["OccluderState"]:
117+
if not vm.occluder_model:
118+
gui.set_status("loading Occluder.")
119+
vm.occluder_model, vm.occluder_tensor = load_occluder_model()
120+
gui.set_status("Occluder loaded!")
121+
if action[0][1]["FaceParserState"]:
122+
if not vm.face_parsing_model:
123+
gui.set_status("loading FaceParser")
124+
vm.face_parsing_model, vm.face_parsing_tensor = load_face_parser_model()
125+
gui.set_status("FaceParser loaded!")
126+
127+
128+
vm.parameters = action[0][1]
129+
action.pop(0)
130+
131+
elif action [0][0] == "load_models":
132+
gui.set_status("loading Faceapp...")
133+
faceapp = load_faceapp_model()
134+
gui.set_faceapp_model(faceapp)
135+
vm.set_faceapp_model(faceapp)
136+
gui.set_status("loading Target Videos...")
137+
gui.populate_target_videos()
138+
gui.set_status("loading Source Faces...")
139+
gui.load_source_faces()
140+
gui.set_status("Done...")
141+
action.pop(0)
142+
143+
144+
# From VM
145+
elif action[0][0] == "stop_play":
146+
gui.toggle_play_video()
147+
action.pop(0)
148+
149+
elif action[0][0] == "set_slider_length":
150+
gui.set_video_slider_length(action[0][1])
151+
action.pop(0)
152+
153+
elif action[0][0] == "send_msg":
154+
gui.set_status(action[0][1])
155+
action.pop(0)
156+
157+
else:
158+
print("Action not found: "+action[0][0]+" "+str(action[0][1]))
159+
action.pop(0)
160+
161+
# start = time.time()
162+
163+
164+
gui.check_for_video_resize()
165+
vm.process()
166+
gui.after(1, coordinator)
167+
# print(time.time() - start)
168+
169+
def load_faceapp_model():
170+
app = FaceAnalysis(name='buffalo_l')
171+
app.prepare(ctx_id=0, det_size=(640, 640))
172+
return app
173+
174+
def load_swapper_model():
175+
# Load Swapper model and get graph param
176+
model = onnx.load("./models/inswapper_128.fp16.onnx")
177+
graph = model.graph
178+
179+
emap = onnx.numpy_helper.to_array(graph.initializer[-1])
180+
181+
# Create Swapper model session
182+
opts = onnxruntime.SessionOptions()
183+
# opts.enable_profiling = True
184+
opts.enable_cpu_mem_arena = False
185+
return onnxruntime.InferenceSession( "./models/inswapper_128.fp16.onnx", opts, providers=["CUDAExecutionProvider"]), emap
186+
187+
def load_clip_model():
188+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
189+
print(f"Using device: {device}")
190+
clip_session = CLIPDensePredT(version='ViT-B/16', reduce_dim=64, complex_trans_conv=True)
191+
clip_session.eval();
192+
clip_session.load_state_dict(torch.load('./models/rd64-uni-refined.pth', map_location=torch.device('cuda')), strict=False)
193+
clip_session.to(device)
194+
return clip_session, device
195+
196+
def load_GFPGAN_model():
197+
198+
opts = onnxruntime.SessionOptions()
199+
# opts.enable_profiling = True
200+
opts.enable_cpu_mem_arena = False
201+
GFPGAN_session = onnxruntime.InferenceSession( "./models/GFPGANv1.4.onnx", providers=["CUDAExecutionProvider"])
202+
return GFPGAN_session
203+
204+
def load_occluder_model():
205+
to_tensor = transforms.ToTensor()
206+
model = smp.Unet(encoder_name='resnet18', encoder_weights='imagenet', classes=1, activation=None)
207+
208+
weights = torch.load('./models/occluder.ckpt')
209+
new_weights = OrderedDict()
210+
for key in weights.keys():
211+
new_key = '.'.join(key.split('.')[1:])
212+
new_weights[new_key] = weights[key]
213+
214+
model.load_state_dict(new_weights)
215+
model.to('cuda')
216+
model.eval()
217+
return model, to_tensor
218+
219+
def load_face_parser_model():
220+
n_classes = 19
221+
model = BiSeNet(n_classes=n_classes)
222+
model.cuda()
223+
model.load_state_dict(torch.load("./models/79999_iter.pth"))
224+
model.eval()
225+
226+
to_tensor = transforms.Compose([
227+
transforms.ToTensor(),
228+
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
229+
])
230+
231+
return model, to_tensor
232+
233+
def run():
234+
global gui, vm, action, frame, r_frame
235+
gui = GUI.GUI()
236+
vm = VM.VideoManager()
237+
238+
action = []
239+
frame = []
240+
r_frame = []
241+
242+
243+
gui.initialize_gui()
244+
coordinator()
245+
246+
gui.mainloop()
247+
248+

0 commit comments

Comments
 (0)