-
Notifications
You must be signed in to change notification settings - Fork 46
/
Copy pathfeat_cluster.py
130 lines (107 loc) · 4.66 KB
/
feat_cluster.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
import argparse
import random
import torch
import torch.nn as nn
from torchvision import utils
from tqdm import tqdm
import sys
import lpips
from torchvision import transforms, utils
from torch.utils import data
import os
from PIL import Image
import numpy as np
def intra_cluster_dist(args):
with torch.no_grad():
lpips_fn = lpips.LPIPS(net='vgg').to(device)
preprocess = transforms.Compose([
transforms.Resize([256, 256]),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
cluster_size = 50
base_path = os.path.join("cluster_centers", "%s" %(args.dataset), "%s" % (args.baseline))
avg_dist = torch.zeros([10, ])
for k in range(10):
curr_path = os.path.join(base_path, "c%d" % (k))
files_list = os.listdir(curr_path)
files_list.remove('center.png')
random.shuffle(files_list)
files_list = files_list[:cluster_size]
dists = []
for i in range(len(files_list)):
for j in range(i+1, len(files_list)):
input1_path = os.path.join(curr_path, files_list[i])
input2_path = os.path.join(curr_path, files_list[j])
input_image1 = Image.open(input1_path)
input_image2 = Image.open(input2_path)
input_tensor1 = preprocess(input_image1)
input_tensor2 = preprocess(input_image2)
input_tensor1 = input_tensor1.to(device)
input_tensor2 = input_tensor2.to(device)
dist = lpips_fn(input_tensor1, input_tensor2)
dists.append(dist)
dists = torch.tensor(dists)
print ("Cluster %d: Avg. pairwise LPIPS dist: %f/%f" %
(k, dists.mean(), dists.std()))
avg_dist[k] = dists.mean()
print ("Final avg. %f/%f" % (avg_dist[~torch.isnan(avg_dist)].mean(), avg_dist[~torch.isnan(avg_dist)].std()))
def get_close_far_members(args):
with torch.no_grad():
lpips_fn = lpips.LPIPS(net='vgg').to(device)
preprocess = transforms.Compose([
transforms.Resize([256, 256]),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
cluster_size = 50
base_path = os.path.join("cluster_centers", "%s" %(args.dataset), "%s" % (args.baseline))
avg_dist = torch.zeros([10, ])
for k in range(10):
curr_path = os.path.join(base_path, "c%d" % (k))
files_list = os.listdir(curr_path)
files_list.remove('center.png')
random.shuffle(files_list)
files_list = files_list[:cluster_size]
dists = []
min_dist, max_dist = 1000.0, 0.0
min_ind, max_ind = 0, 0
# center image
input1_path = os.path.join(curr_path, 'center.png')
input_image1 = Image.open(input1_path)
input_tensor1 = preprocess(input_image1)
input_tensor1 = input_tensor1.to(device)
for i in range(len(files_list)):
input2_path = os.path.join(curr_path, files_list[i])
input_image2 = Image.open(input2_path)
input_tensor2 = preprocess(input_image2)
input_tensor2 = input_tensor2.to(device)
dist = lpips_fn(input_tensor1, input_tensor2)
if dist <= min_dist:
min_ind = i
min_dist = dist
if dist >= max_dist:
max_ind = i
max_dist = dist
print (min_ind, max_ind)
if len(files_list) > 0:
# saving the closest member
path_closest = os.path.join(curr_path, files_list[min_ind])
new_closest = os.path.join(curr_path, 'closest.png')
os.system("cp %s %s" %(path_closest, new_closest))
# saving the farthest member
path_farthest = os.path.join(curr_path, files_list[max_ind])
new_farthest = os.path.join(curr_path, 'farthest.png')
os.system("cp %s %s" %(path_farthest, new_farthest))
else:
print("no members in cluster %d" %(k))
if __name__ == '__main__':
device = 'cuda'
parser = argparse.ArgumentParser()
parser.add_argument("--baseline", type=str, required=True)
parser.add_argument("--dataset", type=str, required=True)
args = parser.parse_args()
intra_cluster_dist(args)
get_close_far_members(args)