3 changed files with 231 additions and 10 deletions
@ -0,0 +1,212 @@ |
|||
from __future__ import print_function |
|||
|
|||
import os |
|||
import sys |
|||
import time |
|||
import torch |
|||
import torch.optim as optim |
|||
import torch.backends.cudnn as cudnn |
|||
import torch.nn.functional as F |
|||
import argparse |
|||
import random |
|||
import numpy as np |
|||
|
|||
from torchvision import transforms, datasets |
|||
import torchvision.models as models |
|||
|
|||
from sklearn.svm import LinearSVC |
|||
|
|||
from scripts.parser import parser |
|||
import scripts.augmentation as aug |
|||
import pcl.loader |
|||
|
|||
def calculate_ap(rec, prec): |
|||
""" |
|||
Computes the AP under the precision recall curve. |
|||
""" |
|||
rec, prec = rec.reshape(rec.size, 1), prec.reshape(prec.size, 1) |
|||
z, o = np.zeros((1, 1)), np.ones((1, 1)) |
|||
mrec, mpre = np.vstack((z, rec, o)), np.vstack((z, prec, z)) |
|||
for i in range(len(mpre) - 2, -1, -1): |
|||
mpre[i] = max(mpre[i], mpre[i + 1]) |
|||
|
|||
indices = np.where(mrec[1:] != mrec[0:-1])[0] + 1 |
|||
ap = 0 |
|||
for i in indices: |
|||
ap = ap + (mrec[i] - mrec[i - 1]) * mpre[i] |
|||
return ap |
|||
|
|||
|
|||
def get_precision_recall(targets, preds): |
|||
""" |
|||
[P, R, score, ap] = get_precision_recall(targets, preds) |
|||
Input : |
|||
targets : number of occurrences of this class in the ith image |
|||
preds : score for this image |
|||
Output : |
|||
P, R : precision and recall |
|||
score : score which corresponds to the particular precision and recall |
|||
ap : average precision |
|||
""" |
|||
# binarize targets |
|||
targets = np.array(targets > 0, dtype=np.float32) |
|||
tog = np.hstack(( |
|||
targets[:, np.newaxis].astype(np.float64), |
|||
preds[:, np.newaxis].astype(np.float64) |
|||
)) |
|||
ind = np.argsort(preds) |
|||
ind = ind[::-1] |
|||
score = np.array([tog[i, 1] for i in ind]) |
|||
sortcounts = np.array([tog[i, 0] for i in ind]) |
|||
|
|||
tp = sortcounts |
|||
fp = sortcounts.copy() |
|||
for i in range(sortcounts.shape[0]): |
|||
if sortcounts[i] >= 1: |
|||
fp[i] = 0. |
|||
elif sortcounts[i] < 1: |
|||
fp[i] = 1. |
|||
P = np.cumsum(tp) / (np.cumsum(tp) + np.cumsum(fp)) |
|||
numinst = np.sum(targets) |
|||
R = np.cumsum(tp) / numinst |
|||
ap = calculate_ap(R, P) |
|||
return P, R, score, ap |
|||
|
|||
|
|||
def main(): |
|||
args = parser().parse_args() |
|||
|
|||
if not args.seed is None: |
|||
random.seed(args.seed) |
|||
np.random.seed(args.seed) |
|||
|
|||
mean = [0.485, 0.456, 0.406] |
|||
std = [0.229, 0.224, 0.225] |
|||
normalize = transforms.Normalize(mean=mean, std=std) |
|||
transform = transforms.Compose([ |
|||
transforms.Resize(256), |
|||
transforms.CenterCrop(224), |
|||
transforms.ToTensor(), |
|||
normalize, |
|||
]) |
|||
|
|||
eval_augmentation = aug.moco_eval() |
|||
|
|||
pre_train_dir = os.path.join(args.data, 'pre_train') |
|||
eval_dir = os.path.join(args.data, 'eval') |
|||
|
|||
train_dataset = pcl.loader.PreImager(pre_train_dir, eval_augmentation) |
|||
val_dataset = pcl.loader.ImageFolderInstance( |
|||
eval_dir, |
|||
eval_augmentation) |
|||
|
|||
val_loader = torch.utils.data.DataLoader( |
|||
val_dataset, batch_size=args.batch_size, shuffle=False, |
|||
num_workers=args.workers, pin_memory=True) |
|||
|
|||
# create model |
|||
print("=> creating model '{}'".format(args.arch)) |
|||
model = models.__dict__[args.arch](num_classes=128) |
|||
|
|||
# load from pre-trained |
|||
if args.pretrained: |
|||
if os.path.isfile(args.pretrained): |
|||
print("=> loading checkpoint '{}'".format(args.pretrained)) |
|||
checkpoint = torch.load(args.pretrained, map_location="cpu") |
|||
state_dict = checkpoint['state_dict'] |
|||
# rename pre-trained keys |
|||
for k in list(state_dict.keys()): |
|||
if k.startswith('module.encoder_q') and not k.startswith('module.encoder_q.fc'): |
|||
# remove prefix |
|||
state_dict[k[len("module.encoder_q."):]] = state_dict[k] |
|||
# delete renamed or unused k |
|||
del state_dict[k] |
|||
model.load_state_dict(state_dict, strict=False) |
|||
model.fc = torch.nn.Identity() |
|||
print("=> loaded pre-trained model '{}'".format(args.pretrained)) |
|||
else: |
|||
print("=> no checkpoint found at '{}'".format(args.pretrained)) |
|||
|
|||
model.cuda() |
|||
model.eval() |
|||
|
|||
test_feats = [] |
|||
test_labels = [] |
|||
print('==> calculate test features') |
|||
for idx, (images, target) in enumerate(val_loader): |
|||
images = images.cuda(non_blocking=True) |
|||
feat = model(images) |
|||
feat = feat.detach().cpu() |
|||
test_feats.append(feat) |
|||
test_labels.append(target) |
|||
|
|||
test_feats = torch.cat(test_feats, 0).numpy() |
|||
test_labels = torch.cat(test_labels, 0).numpy() |
|||
|
|||
test_feats_norm = np.linalg.norm(test_feats, axis=1) |
|||
test_feats = test_feats / (test_feats_norm + 1e-5)[:, np.newaxis] |
|||
|
|||
result = {} |
|||
|
|||
k_list = ['full'] |
|||
|
|||
for k in k_list: |
|||
cost_list = args.cost.split(',') |
|||
result_k = np.zeros(len(cost_list)) |
|||
for i, cost in enumerate(cost_list): |
|||
cost = float(cost) |
|||
avg_map = [] |
|||
for run in range(args.n_run): |
|||
print(len(train_dataset)) |
|||
|
|||
train_loader = torch.utils.data.DataLoader( |
|||
train_dataset, batch_size=args.batch_size, shuffle=False, |
|||
num_workers=args.workers, pin_memory=True) |
|||
|
|||
train_feats = [] |
|||
train_labels = [] |
|||
print('==> calculate train features') |
|||
for idx, (images, target) in enumerate(train_loader): |
|||
images = images.cuda(non_blocking=True) |
|||
feat = model(images) |
|||
feat = feat.detach() |
|||
|
|||
train_feats.append(feat) |
|||
train_labels.append(target) |
|||
|
|||
train_feats = torch.cat(train_feats, 0).cpu().numpy() |
|||
train_labels = torch.cat(train_labels, 0).cpu().numpy() |
|||
|
|||
train_feats_norm = np.linalg.norm(train_feats, axis=1) |
|||
train_feats = train_feats / (train_feats_norm + 1e-5)[:, np.newaxis] |
|||
|
|||
print('==> training SVM Classifier') |
|||
cls_ap = np.zeros((args.num_class, 1)) |
|||
test_labels[test_labels == 0] = -1 |
|||
train_labels[train_labels == 0] = -1 |
|||
for cls in range(args.num_class): |
|||
clf = LinearSVC( |
|||
C=cost, class_weight={1: 2, -1: 1}, intercept_scaling=1.0, |
|||
penalty='l2', loss='squared_hinge', tol=1e-4, |
|||
dual=True, max_iter=2000, random_state=0) |
|||
clf.fit(train_feats, train_labels[:, cls]) |
|||
|
|||
prediction = clf.decision_function(test_feats) |
|||
P, R, score, ap = get_precision_recall(test_labels[:, cls], prediction) |
|||
cls_ap[cls][0] = ap * 100 |
|||
mean_ap = np.mean(cls_ap, axis=0) |
|||
|
|||
print('==> Run%d mAP is %.2f: ' % (run, mean_ap)) |
|||
avg_map.append(mean_ap) |
|||
|
|||
avg_map = np.asarray(avg_map) |
|||
print('Cost:%.2f - Average ap is: %.2f' % (cost, avg_map.mean())) |
|||
print('Cost:%.2f - Std is: %.2f' % (cost, avg_map.std())) |
|||
result_k[i] = avg_map.mean() |
|||
result[k] = result_k.max() |
|||
print(result) |
|||
|
|||
|
|||
if __name__ == '__main__': |
|||
main() |
|||
|
|||
Loading…
Reference in new issue