|
4 | 4 | import torch
|
5 | 5 | from torchvision import transforms
|
6 | 6 | from .nets import DSFDNet
|
7 |
| -from .nets2 import DSFDNet2 |
8 | 7 | from .box_utils import nms_
|
9 | 8 |
|
10 | 9 | PATH_WEIGHT = './detectors/dsfd/weights/dsfd_vgg_0.880.pth'
|
11 |
| -PATH_WEIGHT2 = './detectors/dsfd/weights/WIDERFace_DSFD_RES152.pth' |
12 | 10 | img_mean = np.array([104., 117., 123.])[:, np.newaxis, np.newaxis].astype('float32')
|
13 | 11 |
|
14 | 12 |
|
@@ -61,51 +59,3 @@ def detect_faces(self, image, conf_th=0.8, scales=[1]):
|
61 | 59 | bboxes = bboxes[keep]
|
62 | 60 |
|
63 | 61 | return bboxes
|
64 |
| - |
65 |
| - |
66 |
| -class DSFD_(): |
67 |
| - |
68 |
| - def __init__(self, device='cuda'): |
69 |
| - |
70 |
| - tstamp = time.time() |
71 |
| - self.device = device |
72 |
| - |
73 |
| - print('[DSFD] loading with', self.device) |
74 |
| - self.net = DSFDNet2(device=self.device).to(self.device) |
75 |
| - state_dict = torch.load(PATH_WEIGHT2, map_location=self.device) |
76 |
| - self.net.load_state_dict(state_dict) |
77 |
| - self.net.eval() |
78 |
| - print('[DSFD] finished loading (%.4f sec)' % (time.time() - tstamp)) |
79 |
| - |
80 |
| - def detect_faces(self, image, conf_th=0.8, scales=[1]): |
81 |
| - |
82 |
| - w, h = image.shape[1], image.shape[0] |
83 |
| - |
84 |
| - bboxes = np.empty(shape=(0, 5)) |
85 |
| - |
86 |
| - with torch.no_grad(): |
87 |
| - for s in scales: |
88 |
| - scaled_img = cv2.resize(image, dsize=(0, 0), fx=s, fy=s, interpolation=cv2.INTER_LINEAR) |
89 |
| - |
90 |
| - scaled_img = np.swapaxes(scaled_img, 1, 2) |
91 |
| - scaled_img = np.swapaxes(scaled_img, 1, 0) |
92 |
| - scaled_img = scaled_img[[2, 1, 0], :, :] |
93 |
| - scaled_img = scaled_img.astype('float32') |
94 |
| - scaled_img -= img_mean |
95 |
| - scaled_img = scaled_img[[2, 1, 0], :, :] |
96 |
| - x = torch.from_numpy(scaled_img).unsqueeze(0).to(self.device) |
97 |
| - y = self.net(x) |
98 |
| - |
99 |
| - detections = y.data |
100 |
| - scale = torch.Tensor([w, h, w, h]) |
101 |
| - |
102 |
| - for i in range(detections.size(1)): |
103 |
| - j = 0 |
104 |
| - while detections[0, i, j, 0] > conf_th: |
105 |
| - score = detections[0, i, j, 0] |
106 |
| - pt = (detections[0, i, j, 1:] * scale).cpu().numpy() |
107 |
| - bbox = (pt[0], pt[1], pt[2], pt[3], score) |
108 |
| - bboxes = np.vstack((bboxes, bbox)) |
109 |
| - j += 1 |
110 |
| - |
111 |
| - return bboxes |
0 commit comments