-
Notifications
You must be signed in to change notification settings - Fork 24
/
Copy pathdavis_eval.py
79 lines (76 loc) · 3.03 KB
/
davis_eval.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
import os
import sys
import cv2
from PIL import Image
import numpy as np
from util import calcIoU
PALETTE = [0, 0, 0, 128, 0, 0, 0, 128, 0, 128, 128, 0, 0, 0, 128, 128, 0, 128, 0, 128, 128, 128, 128, 128, 64, 0, 0, 191, 0, 0, 64, 128, 0, 191, 128, 0, 64, 0, 128]
data_path = sys.argv[1]
pred_path = sys.argv[2] #'DAVIS/Results/Segmentations/480p/OSVOS'
dataset_version = sys.argv[3]
dataset_split = sys.argv[4]
if len(sys.argv) > 5:
vis_path = sys.argv[5]
else:
vis_path = None
listFile = '%s/ImageSets/%s/%s.txt' % (data_path, dataset_version, dataset_split)
gt_path = os.path.join(data_path, 'Annotations', '480p')
with open(listFile, 'r') as f:
fds = [line.strip() for line in f]
im_num = 0
iou =[]
seq_n = 0
sample_n = 0
subfd_names = []
for i, fd in enumerate(fds):
print fd
file_list = os.listdir(os.path.join(gt_path, fd))
im_list = [name for name in file_list if len(name) > 4 and name[-4:]=='.png']
im_list = sorted(im_list)
im_list = im_list[1:-1] # remove first and last image
pred_list = os.listdir(os.path.join(pred_path, fd))
if dataset_version == '2017':
sub_fds = [name for name in pred_list if len(name) < 4]
sub_fds = sorted(sub_fds)
print sub_fds
for sub_fd in sub_fds:
subfd_names.append(fd+'/'+sub_fd)
iou_seq = []
for i,im_name in enumerate(im_list):
iou_im = 0
scores = []
label_gt = np.array(Image.open(os.path.join(gt_path, fd, im_name)))
if dataset_version == '2017':
for j, sub_fd in enumerate(sub_fds):
score = np.load(os.path.join(pred_path, fd, sub_fd, im_name[:-4] + '.npy'))
scores.append(score)
im_size = scores[0].shape
bg_score = np.ones(im_size) * 0.5
scores = [bg_score] + scores
score_all = np.stack(tuple(scores), axis = -1)
class_n = score_all.shape[2] - 1
label_pred = score_all.argmax(axis=2)
else:
class_n = 1
label_gt = label_gt > 0
label_pred = np.array(Image.open(os.path.join(pred_path,fd, im_name))) > 0
label_pred = np.array(Image.fromarray(label_pred.astype(np.uint8)).resize((label_gt.shape[1], label_gt.shape[0]),
Image.NEAREST))
#cv2.resize(label_pred, label_gt.shape, label_pred, 0, 0, cv2.INTER_NEAREST)
if vis_path:
res_im = Image.fromarray(label_pred, mode="P")
res_im.putpalette(PALETTE)
if not os.path.exists(os.path.join(vis_path, fd)):
os.makedirs(os.path.join(vis_path, fd))
res_im.save(os.path.join(vis_path, fd, im_name))
iou_seq.append(calcIoU(label_gt, label_pred, class_n))
iou_seq = np.stack(iou_seq, axis=1)
print iou_seq.mean(axis=1)
sample_n += iou_seq.size
iou.extend(iou_seq.mean(axis=1).tolist())#flatten and append
iou = np.array(iou)
print "iou:", iou.mean()
with open("iou.txt", "w") as f:
for fd, num in zip(subfd_names, iou):
f.write("%s\t%f\n" % (fd, num))
f.write("all\t%f\n" % iou.mean())