conciomith commited on
Commit
9e6296d
1 Parent(s): 7974733

Upload preprocess.py

Browse files
Files changed (1) hide show
  1. preprocess.py +53 -0
preprocess.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import cv2
3
+
4
+
5
+ # This function is modified from the following code snippet:
6
+ # https://github.com/StanislasBertrand/RetinaFace-tf2/blob/5f68ce8130889384cb8aca937a270cea4ef2d020/retinaface.py#L49-L74
7
+ def resize_image(img, scales, allow_upscaling):
8
+ img_h, img_w = img.shape[0:2]
9
+ target_size = scales[0]
10
+ max_size = scales[1]
11
+
12
+ if img_w > img_h:
13
+ im_size_min, im_size_max = img_h, img_w
14
+ else:
15
+ im_size_min, im_size_max = img_w, img_h
16
+
17
+ im_scale = target_size / float(im_size_min)
18
+ if not allow_upscaling:
19
+ im_scale = min(1.0, im_scale)
20
+
21
+ if np.round(im_scale * im_size_max) > max_size:
22
+ im_scale = max_size / float(im_size_max)
23
+
24
+ if im_scale != 1.0:
25
+ img = cv2.resize(
26
+ img,
27
+ None,
28
+ None,
29
+ fx=im_scale,
30
+ fy=im_scale,
31
+ interpolation=cv2.INTER_LINEAR
32
+ )
33
+
34
+ return img, im_scale
35
+
36
+
37
+ # This function is modified from the following code snippet:
38
+ # https://github.com/StanislasBertrand/RetinaFace-tf2/blob/5f68ce8130889384cb8aca937a270cea4ef2d020/retinaface.py#L76-L96
39
+ def preprocess_image(img, allow_upscaling):
40
+ pixel_means = np.array([0.0, 0.0, 0.0], dtype=np.float32)
41
+ pixel_stds = np.array([1.0, 1.0, 1.0], dtype=np.float32)
42
+ pixel_scale = float(1.0)
43
+ scales = [1024, 1980]
44
+
45
+ img, im_scale = resize_image(img, scales, allow_upscaling)
46
+ img = img.astype(np.float32)
47
+ im_tensor = np.zeros((1, img.shape[0], img.shape[1], img.shape[2]), dtype=np.float32)
48
+
49
+ # Make image scaling + BGR2RGB conversion + transpose (N,H,W,C) to (N,C,H,W)
50
+ for i in range(3):
51
+ im_tensor[0, :, :, i] = (img[:, :, 2 - i] / pixel_scale - pixel_means[2 - i]) / pixel_stds[2 - i]
52
+
53
+ return im_tensor, img.shape[0:2], im_scale