-
Notifications
You must be signed in to change notification settings - Fork 0
/
webcam_face_sr.py
165 lines (124 loc) · 6.12 KB
/
webcam_face_sr.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
import numpy as np
import cv2
import os
import sys
import matplotlib.pyplot as plt
import torch
from PIL import Image
from dlib_alignment import dlib_detect_face, face_recover
import torchvision.transforms as transforms
from models.SRGAN_model import SRGANModel
import argparse
import utils
def get_FaceSR_opt():
parser = argparse.ArgumentParser()
parser.add_argument('--gpu_ids', type=str, default=None)
parser.add_argument('--batch_size', type=int, default=32)
parser.add_argument('--lr_G', type=float, default=1e-4)
parser.add_argument('--weight_decay_G', type=float, default=0)
parser.add_argument('--beta1_G', type=float, default=0.9)
parser.add_argument('--beta2_G', type=float, default=0.99)
parser.add_argument('--lr_D', type=float, default=1e-4)
parser.add_argument('--weight_decay_D', type=float, default=0)
parser.add_argument('--beta1_D', type=float, default=0.9)
parser.add_argument('--beta2_D', type=float, default=0.99)
parser.add_argument('--lr_scheme', type=str, default='MultiStepLR')
parser.add_argument('--niter', type=int, default=100000)
parser.add_argument('--warmup_iter', type=int, default=-1)
parser.add_argument('--lr_steps', type=list, default=[50000])
parser.add_argument('--lr_gamma', type=float, default=0.5)
parser.add_argument('--pixel_criterion', type=str, default='l1')
parser.add_argument('--pixel_weight', type=float, default=1e-2)
parser.add_argument('--feature_criterion', type=str, default='l1')
parser.add_argument('--feature_weight', type=float, default=1)
parser.add_argument('--gan_type', type=str, default='ragan')
parser.add_argument('--gan_weight', type=float, default=5e-3)
parser.add_argument('--D_update_ratio', type=int, default=1)
parser.add_argument('--D_init_iters', type=int, default=0)
parser.add_argument('--print_freq', type=int, default=100)
parser.add_argument('--val_freq', type=int, default=1000)
parser.add_argument('--save_freq', type=int, default=10000)
parser.add_argument('--crop_size', type=float, default=0.85)
parser.add_argument('--lr_size', type=int, default=128)
parser.add_argument('--hr_size', type=int, default=512)
# network G
parser.add_argument('--which_model_G', type=str, default='RRDBNet')
parser.add_argument('--G_in_nc', type=int, default=3)
parser.add_argument('--out_nc', type=int, default=3)
parser.add_argument('--G_nf', type=int, default=64)
parser.add_argument('--nb', type=int, default=16)
# network D
parser.add_argument('--which_model_D', type=str, default='discriminator_vgg_128')
parser.add_argument('--D_in_nc', type=int, default=3)
parser.add_argument('--D_nf', type=int, default=64)
# data dir
parser.add_argument('--pretrain_model_G', type=str, default='90000_G.pth')
parser.add_argument('--pretrain_model_D', type=str, default=None)
args = parser.parse_args("")
return args
border = 0.25
def set_res(cap, x,y):
cap.set(3, int(x))
cap.set(4, int(y))
return True
def main():
cap = cv2.VideoCapture(0)
print(set_res(cap, 640, 480))
cv2.namedWindow('SR', cv2.WINDOW_NORMAL)
# setting up the model
_transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])])
sr_model = SRGANModel(get_FaceSR_opt(), is_train=False)
sr_model.load()
def sr_forward(img, padding=0.5, moving=0.1):
# img_aligned, M = dlib_detect_face(img, padding=padding, image_size=img.shape[:-1], moving=moving)
input_img = torch.unsqueeze(_transform(Image.fromarray(img)), 0)
sr_model.var_L = input_img.to(sr_model.device)
sr_model.test()
output_img = sr_model.fake_H.squeeze(0).cpu().numpy()
output_img = np.clip((np.transpose(output_img, (1, 2, 0)) / 2.0 + 0.5) * 255.0, 0, 255).astype(np.uint8)
# rec_img = face_recover(output_img, M * 4, img)
return output_img
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
def single_image_face_sr(image):
# getting faces
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
faces = face_cascade.detectMultiScale(gray, 1.1, 4)
print("Faces:", len(faces))
face_crops = []
for x, y, w, h in faces:
offsets = (int(w * border), int(h * border)) # w, h
point = (max(0, x - offsets[0]), max(0, y - offsets[1]))
face_crop_bgr = image[point[1]: point[1] + h + 2 * offsets[1], \
point[0]: point[0] + w + 2 * offsets[0]]
face_crop = cv2.cvtColor(face_crop_bgr, cv2.COLOR_BGR2RGB)
print(face_crop.shape)
face_crops.append((point, face_crop))
new_image_bgr = cv2.resize(image, (image.shape[1] * 4, image.shape[0] * 4))
new_image = cv2.cvtColor(new_image_bgr, cv2.COLOR_BGR2RGB)
for face in face_crops:
face_image = face[1]
position = face[0] + (face_image.shape)[:-1]
output_img = sr_forward(face_image)
new_dims = output_img.shape[:-1]
new_image[4 * position[1]: 4 * position[1] + new_dims[1], 4 * position[0] : 4 * position[0] + new_dims[0]] = output_img
new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
return new_image
# setting up face detector
while True:
ret, img = cap.read()
# img = cv2.flip(img, -1)
# the model x4 the faces, but runs too slowly on the full 1080p output, so
# I 'cheat' and half the resolution, therefore the faces are still x2 the
# resolution as normal in order to run even close to real-time.
img = cv2.resize(img, (int(img.shape[1] / 2), int(img.shape[0] / 2)))
print(img.shape)
img_sr = single_image_face_sr(img)
cv2.imshow('video', img_sr)
k = cv2.waitKey(30) & 0xff
if k == 27: # press 'ESC' to quit
break
cap.release()
cv2.destroyAllWindows()
main()