-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
158 lines (126 loc) · 4.9 KB
/
utils.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
import sys
import os
import numpy
import cv2
import os
import json
import math
from PIL import Image
from matplotlib import pyplot as plt
import torch
import torchvision
from torchvision import datasets, transforms
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import _LRScheduler
import conf
def get_model(model_type, use_gpu):
if model_type == 'vgg16':
from models.vgg import vgg16
model = vgg16()
elif model_type == 'resnet50':
from models.resnet import resnet50
model = resnet50()
elif model_type == 'resnet18':
from models.resnet import resnet18
model = resnet18()
elif model_type == 'googlenet':
from models.googlenet import googlenet
model = googlenet()
else:
print('this model is not supported')
sys.exit()
if use_gpu:
model = model.cuda()
return model
def my_eval(model, data_path, use_gpu):
model.eval()
final_dict = {}
json_dict = {}
for filename in os.listdir(data_path):
if filename.endswith('.mp4') or filename.endswith('.avi'):
#continue
video_path = data_path + '/' + filename
#print(video_path)
reader = cv2.VideoCapture(video_path)
num_frames = int(reader.get(cv2.CAP_PROP_FRAME_COUNT))
ans = 0
real_ans = 0
fake_ans = 0
all_poss = 0.0
while reader.isOpened():
_, image = reader.read()
if image is None:
break
ans += 1
if ans % conf.FRAME_SAMPLE != 0:
continue
image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
image = image.resize((conf.IMAGE_SIZE, conf.IMAGE_SIZE),Image.ANTIALIAS)
transform = transforms.Compose([transforms.ToTensor()])
img = transform(image)
img = img.resize(1, 3, conf.IMAGE_SIZE, conf.IMAGE_SIZE)
if use_gpu:
output = model(img.cuda())
else:
output = model(img)
all_poss += math.exp(output[0][0].cpu().detach().numpy()) / (math.exp(output[0][0].cpu().detach().numpy()) + math.exp(output[0][1].cpu().detach().numpy()))
_, pred = output.topk(1, 1, True, True)
if pred.cpu().numpy()[0][0] == 1:
real_ans += 1
else:
fake_ans += 1
if fake_ans + real_ans != 0:
fake_poss = float(all_poss) / (fake_ans + real_ans)
else:
fake_poss = 0.5
final_dict[filename] = fake_poss
print(filename + ': ' + str(fake_poss))
elif filename.endswith('.json'):
print('found json')
json_file = open(data_path + '/' + filename, encoding='utf-8')
json_dict = json.load(json_file)
# check
correct_ans = 0
score = 0
for key in final_dict:
if key in json_dict:
if json_dict[key]['label'] == 'FAKE' if final_dict[key] > 0.5 else 'REAL':
correct_ans += 1
yi = 1 if json_dict[key]['label'] == 'FAKE' else 0
if final_dict[key] >= 1:
final_dict[key] = 0.999
elif final_dict[key] <= 0:
final_dict[key] = 0.001
score += yi * math.log(final_dict[key]) + (1 - yi) * math.log(1 - final_dict[key])
else:
print('dict key unmatched')
score = - score / len(final_dict)
print('final accuracy:' + str(float(correct_ans) / len(final_dict)))
print('final score: ' + str(score))
def resize_image(image, resize_height = 256, resize_width = 256):
image_shape = numpy.shape(image)
height = image_shape[0]
width = image_shape[1]
image = cv2.resize(image, dsize=(resize_width, resize_height))
return image
def image_preprocess(filename, resize_height = 256, resize_width = 256, normalization=False):
bgr_image = cv2.imread(filename)
if bgr_image is None:
print("Image does not exist: ", filename)
return None
if len(bgr_image.shape) == 2:
print("Warning: gray image: ", filename)
bgr_image = cv2.cvtColor(bgr_image, cv2.COLOR_GRAY2BGR)
rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
rgb_image = resize_image(rgb_image, resize_height, resize_width)
rgb_image = numpy.asanyarray(rgb_image)
if normalization:
rgb_image = rgb_image / 255.0
return rgb_image
class WarmUpLR(_LRScheduler):
def __init__(self, optimizer, total_iters, last_epoch=-1):
self.total_iters = total_iters
super().__init__(optimizer, last_epoch)
def get_lr(self):
return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]