Skip to content

Commit 93cba35

Browse files
DanyacheDanyache
committed
added inference pipeline
1 parent 0aff2cc commit 93cba35

8 files changed

+381
-12
lines changed

‎SberSwapInference.ipynb

+369
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,369 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": 11,
6+
"id": "mental-laser",
7+
"metadata": {},
8+
"outputs": [
9+
{
10+
"name": "stderr",
11+
"output_type": "stream",
12+
"text": [
13+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
14+
" and should_run_async(code)\n"
15+
]
16+
}
17+
],
18+
"source": [
19+
"import cv2\n",
20+
"import torch\n",
21+
"import time\n",
22+
"\n",
23+
"from utils.image_processing import crop_face\n",
24+
"from utils.video_processing import read_video, get_final_video_frame, add_audio_from_another_video\n",
25+
"from utils.core import model_inference\n",
26+
"\n",
27+
"from network.AEI_Net import AEI_Net\n",
28+
"from coordinate_reg.image_infer import Handler\n",
29+
"from insightface_func.face_detect_crop_single import Face_detect_crop\n",
30+
"from arcface_model.iresnet import iresnet100"
31+
]
32+
},
33+
{
34+
"cell_type": "markdown",
35+
"id": "green-navigation",
36+
"metadata": {},
37+
"source": [
38+
"### Load Models"
39+
]
40+
},
41+
{
42+
"cell_type": "code",
43+
"execution_count": 12,
44+
"id": "reported-optimum",
45+
"metadata": {},
46+
"outputs": [
47+
{
48+
"name": "stderr",
49+
"output_type": "stream",
50+
"text": [
51+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
52+
" and should_run_async(code)\n"
53+
]
54+
},
55+
{
56+
"name": "stdout",
57+
"output_type": "stream",
58+
"text": [
59+
"input mean and std: 127.5 127.5\n",
60+
"find model: ./insightface_func/models/antelope/glintr100.onnx recognition\n",
61+
"find model: ./insightface_func/models/antelope/scrfd_10g_bnkps.onnx detection\n",
62+
"set det-size: (640, 640)\n",
63+
"loading ./coordinate_reg/model/2d106det 0\n",
64+
"input mean and std: 127.5 127.5\n",
65+
"find model: ./insightface_func/models/antelope/glintr100.onnx recognition\n",
66+
"find model: ./insightface_func/models/antelope/scrfd_10g_bnkps.onnx detection\n",
67+
"set det-size: (640, 640)\n"
68+
]
69+
}
70+
],
71+
"source": [
72+
"app = Face_detect_crop(name='antelope', root='./insightface_func/models')\n",
73+
"app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))\n",
74+
"\n",
75+
"# основная модель для генерации\n",
76+
"G = AEI_Net(c_id=512)\n",
77+
"G.eval()\n",
78+
"G.load_state_dict(torch.load('weights/G_0_035000_init_arch_arcface2.pth', map_location=torch.device('cpu')))\n",
79+
"G = G.cuda()\n",
80+
"G = G.half()\n",
81+
"\n",
82+
"# модель arcface для того, чтобы достать эмбеддинг лица\n",
83+
"netArc = iresnet100(fp16=False)\n",
84+
"netArc.load_state_dict(torch.load('arcface_model/backbone.pth'))\n",
85+
"netArc=netArc.cuda()\n",
86+
"netArc.eval()\n",
87+
"\n",
88+
"# модель, которая позволяет находить точки лица\n",
89+
"handler = Handler('./coordinate_reg/model/2d106det', 0, ctx_id=0, det_size=640)"
90+
]
91+
},
92+
{
93+
"cell_type": "markdown",
94+
"id": "studied-display",
95+
"metadata": {},
96+
"source": [
97+
"### Set here path to source image and video for faceswap"
98+
]
99+
},
100+
{
101+
"cell_type": "code",
102+
"execution_count": 13,
103+
"id": "governmental-anthropology",
104+
"metadata": {},
105+
"outputs": [
106+
{
107+
"name": "stderr",
108+
"output_type": "stream",
109+
"text": [
110+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
111+
" and should_run_async(code)\n"
112+
]
113+
}
114+
],
115+
"source": [
116+
"\"\"\"\n",
117+
"choose not really long videos, coz it can take a lot of time othervise \n",
118+
"choose source image as a photo -- preferable a selfie of a person\n",
119+
"\"\"\"\n",
120+
"\n",
121+
"path_to_video = 'examples/videos/dora_short.mp4'\n",
122+
"source_full = cv2.imread('examples/images/elon_musk.jpg')\n",
123+
"OUT_VIDEO_NAME = \"examples/results/elon2dora.mp4\"\n",
124+
"crop_size = 224 # don't change this"
125+
]
126+
},
127+
{
128+
"cell_type": "code",
129+
"execution_count": 14,
130+
"id": "forty-tomorrow",
131+
"metadata": {},
132+
"outputs": [
133+
{
134+
"name": "stdout",
135+
"output_type": "stream",
136+
"text": [
137+
"Everything is ok!\n"
138+
]
139+
},
140+
{
141+
"name": "stderr",
142+
"output_type": "stream",
143+
"text": [
144+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
145+
" and should_run_async(code)\n"
146+
]
147+
}
148+
],
149+
"source": [
150+
"# check, if we can detect face on the source image\n",
151+
"\n",
152+
"try:\n",
153+
" source = crop_face(source_full, app, crop_size)[0]\n",
154+
" source = source[:, :, ::-1]\n",
155+
" print(\"Everything is ok!\")\n",
156+
"except TypeError:\n",
157+
" print(\"Bad source image. Choose another one.\")"
158+
]
159+
},
160+
{
161+
"cell_type": "code",
162+
"execution_count": 15,
163+
"id": "third-reset",
164+
"metadata": {},
165+
"outputs": [
166+
{
167+
"name": "stderr",
168+
"output_type": "stream",
169+
"text": [
170+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
171+
" and should_run_async(code)\n"
172+
]
173+
}
174+
],
175+
"source": [
176+
"# read video\n",
177+
"\n",
178+
"full_frames, fps = read_video(path_to_video)"
179+
]
180+
},
181+
{
182+
"cell_type": "markdown",
183+
"id": "reverse-supervision",
184+
"metadata": {},
185+
"source": [
186+
"### Model Inference"
187+
]
188+
},
189+
{
190+
"cell_type": "code",
191+
"execution_count": 16,
192+
"id": "moved-process",
193+
"metadata": {},
194+
"outputs": [
195+
{
196+
"name": "stderr",
197+
"output_type": "stream",
198+
"text": [
199+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
200+
" and should_run_async(code)\n"
201+
]
202+
}
203+
],
204+
"source": [
205+
"START_TIME = time.time()"
206+
]
207+
},
208+
{
209+
"cell_type": "code",
210+
"execution_count": 17,
211+
"id": "coordinated-terminology",
212+
"metadata": {},
213+
"outputs": [
214+
{
215+
"name": "stderr",
216+
"output_type": "stream",
217+
"text": [
218+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
219+
" and should_run_async(code)\n",
220+
"100%|██████████| 262/262 [00:02<00:00, 102.10it/s]\n",
221+
"262it [00:00, 6269.08it/s]\n",
222+
"100%|██████████| 5/5 [00:01<00:00, 3.92it/s]\n",
223+
"100%|██████████| 262/262 [00:00<00:00, 914993.88it/s]\n"
224+
]
225+
}
226+
],
227+
"source": [
228+
"final_frames, crop_frames, full_frames, tfm_array = model_inference(full_frames,\n",
229+
" source,\n",
230+
" [netArc],\n",
231+
" G,\n",
232+
" app, \n",
233+
" crop_size=crop_size)"
234+
]
235+
},
236+
{
237+
"cell_type": "code",
238+
"execution_count": 18,
239+
"id": "spatial-lambda",
240+
"metadata": {},
241+
"outputs": [
242+
{
243+
"name": "stderr",
244+
"output_type": "stream",
245+
"text": [
246+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
247+
" and should_run_async(code)\n",
248+
"100%|██████████| 262/262 [00:08<00:00, 30.33it/s]\n"
249+
]
250+
}
251+
],
252+
"source": [
253+
"get_final_video_frame(final_frames,\n",
254+
" crop_frames,\n",
255+
" full_frames,\n",
256+
" tfm_array,\n",
257+
" OUT_VIDEO_NAME,\n",
258+
" fps, \n",
259+
" handler)"
260+
]
261+
},
262+
{
263+
"cell_type": "code",
264+
"execution_count": 19,
265+
"id": "celtic-press",
266+
"metadata": {},
267+
"outputs": [
268+
{
269+
"name": "stderr",
270+
"output_type": "stream",
271+
"text": [
272+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
273+
" and should_run_async(code)\n"
274+
]
275+
}
276+
],
277+
"source": [
278+
"add_audio_from_another_video(path_to_video, OUT_VIDEO_NAME, \"audio\")"
279+
]
280+
},
281+
{
282+
"cell_type": "code",
283+
"execution_count": 20,
284+
"id": "artificial-mirror",
285+
"metadata": {},
286+
"outputs": [
287+
{
288+
"name": "stdout",
289+
"output_type": "stream",
290+
"text": [
291+
"Full pipeline took 16.079718351364136\n"
292+
]
293+
},
294+
{
295+
"name": "stderr",
296+
"output_type": "stream",
297+
"text": [
298+
"/home/jovyan/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
299+
" and should_run_async(code)\n"
300+
]
301+
}
302+
],
303+
"source": [
304+
"print(f'Full pipeline took {time.time() - START_TIME}')"
305+
]
306+
},
307+
{
308+
"cell_type": "code",
309+
"execution_count": null,
310+
"id": "rotary-compromise",
311+
"metadata": {},
312+
"outputs": [],
313+
"source": []
314+
},
315+
{
316+
"cell_type": "code",
317+
"execution_count": null,
318+
"id": "arbitrary-setup",
319+
"metadata": {},
320+
"outputs": [],
321+
"source": []
322+
},
323+
{
324+
"cell_type": "code",
325+
"execution_count": null,
326+
"id": "offshore-narrow",
327+
"metadata": {},
328+
"outputs": [],
329+
"source": []
330+
},
331+
{
332+
"cell_type": "code",
333+
"execution_count": null,
334+
"id": "grave-classics",
335+
"metadata": {},
336+
"outputs": [],
337+
"source": []
338+
},
339+
{
340+
"cell_type": "code",
341+
"execution_count": null,
342+
"id": "cutting-mumbai",
343+
"metadata": {},
344+
"outputs": [],
345+
"source": []
346+
}
347+
],
348+
"metadata": {
349+
"kernelspec": {
350+
"display_name": "facesh",
351+
"language": "python",
352+
"name": "facesh"
353+
},
354+
"language_info": {
355+
"codemirror_mode": {
356+
"name": "ipython",
357+
"version": 3
358+
},
359+
"file_extension": ".py",
360+
"mimetype": "text/x-python",
361+
"name": "python",
362+
"nbconvert_exporter": "python",
363+
"pygments_lexer": "ipython3",
364+
"version": "3.8.8"
365+
}
366+
},
367+
"nbformat": 4,
368+
"nbformat_minor": 5
369+
}

‎examples/images/elon_musk.jpg

168 KB
Loading

‎examples/results/elon2dora.mp4

9.38 MB
Binary file not shown.

‎examples/videos/dora_short.mp4

10.1 MB
Binary file not shown.

‎examples/videos/efiop_short.mp4

6.9 MB
Binary file not shown.

‎examples/videos/obama_short.mp4

2.25 MB
Binary file not shown.

‎examples/videos/random_gif.gif

1.34 MB
Loading

0 commit comments

Comments
 (0)