forked from clcarwin/SFD_pytorch
-
Notifications
You must be signed in to change notification settings - Fork 1
/
wider_eval_pytorch.py
106 lines (87 loc) · 3.34 KB
/
wider_eval_pytorch.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
from __future__ import print_function
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
# torch.backends.cudnn.bencmark = True
import os,sys,cv2,random,datetime,math
import argparse
import numpy as np
import scipy.io as sio
import zipfile
from net_s3fd import s3fd
from bbox import *
from detect_faces import detect
def flip_detect(net,img):
img = cv2.flip(img, 1)
b = detect(net,img)
bboxlist = np.zeros(b.shape)
bboxlist[:, 0] = img.shape[1] - b[:, 2]
bboxlist[:, 1] = b[:, 1]
bboxlist[:, 2] = img.shape[1] - b[:, 0]
bboxlist[:, 3] = b[:, 3]
bboxlist[:, 4] = b[:, 4]
return bboxlist
def scale_detect(net,img,scale=2.0,facesize=None):
img = cv2.resize(img,(0,0),fx=scale,fy=scale)
b = detect(net,img)
bboxlist = np.zeros(b.shape)
bboxlist[:, 0] = b[:, 0]/scale
bboxlist[:, 1] = b[:, 1]/scale
bboxlist[:, 2] = b[:, 2]/scale
bboxlist[:, 3] = b[:, 3]/scale
bboxlist[:, 4] = b[:, 4]
b = bboxlist
if scale>1: index = np.where(np.minimum(b[:,2]-b[:,0]+1,b[:,3]-b[:,1]+1)<facesize)[0] # only detect small face
else: index = np.where(np.maximum(b[:,2]-b[:,0]+1,b[:,3]-b[:,1]+1)>facesize)[0] # only detect large face
bboxlist = b[index,:]
if 0==len(bboxlist): bboxlist=np.zeros((1, 5))
return bboxlist
wider_face_mat = sio.loadmat('./eval/wider_face_val.mat')
event_list = wider_face_mat['event_list']
file_list = wider_face_mat['file_list']
save_path = './eval/sfd_val_pytorch/'
dataset = '../../dataset/face/WIDER/WIDER_val.zip'
datazip = zipfile.ZipFile(dataset)
net = s3fd()
net.load_state_dict(torch.load('data/s3fd_convert.pth'))
net.cuda()
net.eval()
# for i in range(1000):
# size = 1024+64*i; print(size)
# detect(net,np.zeros((size,size,3)))
for index, event in enumerate(event_list):
filelist = file_list[index][0]
im_dir = event[0][0].encode('utf-8')
if not os.path.exists(save_path + im_dir): os.makedirs(save_path + im_dir)
for num, file in enumerate(filelist):
im_name = file[0][0].encode('utf-8')
zipname = '%s/%s.jpg' % (im_dir,im_name)
data = np.frombuffer(datazip.read('WIDER_val/images/'+zipname),np.uint8)
img = cv2.imdecode(data,1)
imgshow = np.copy(img)
b1 = detect(net,img)
b2 = flip_detect(net,img)
if img.shape[0]*img.shape[1]*4>3000*3000: b3 = np.zeros((1, 5))
else: b3 = scale_detect(net,img,scale=2,facesize=100)
b4 = scale_detect(net,img,scale=0.5,facesize=100)
bboxlist = np.concatenate((b1,b2,b3,b4))
keep = nms(bboxlist,0.3)
keep = keep[0:750] # keep only max 750 boxes
bboxlist = bboxlist[keep,:]
# for b in bboxlist:
# x1,y1,x2,y2,s = b
# if s<0.5: continue
# cv2.rectangle(imgshow,(int(x1),int(y1)),(int(x2),int(y2)),(0,0,255),1)
# cv2.imshow('',imgshow)
# cv2.waitKey(0)
# continue
f = open(save_path + im_dir + '/' + im_name + '.txt', 'w')
f.write('{:s}\n'.format('%s/%s.jpg' % (im_dir,im_name)))
f.write('{:d}\n'.format(len(bboxlist)))
for b in bboxlist:
x1,y1,x2,y2,s = b
f.write('{:.1f} {:.1f} {:.1f} {:.1f} {:.3f}\n'.format(x1,y1,(x2-x1+1),(y2-y1+1),s))
f.close()
print('event:%d num:%d' % (index + 1, num + 1))