-
Notifications
You must be signed in to change notification settings - Fork 1
/
runLAD.py
92 lines (73 loc) · 3.12 KB
/
runLAD.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
import glob
import random
from collie.models import *
random.seed(0)
iterations = 50
rounds = 6
n_labels_teach = 50
n_labels_test = 50
n_labels_tot = 200
# Set this to your LAD image directory
dataDir = "~/data/LAD/images"
allNames = collie.read_text_file("unusualNames.txt")
out = open("result-lad.tsv", "w")
out.write("iter\tround\ttext\tmethod\tcategory\tlabel\ttraining\trank\n")
for iteration in range(iterations):
print(f"Iteration {iteration}")
labels = []
categories = []
imageFiles = []
dirs = []
for dir in glob.glob(dataDir + '/*'):
category = os.path.basename(dir)[0:1]
if category in ["E", "A", "F", "V"]:
dirs.append(dir)
if len(dirs) == 0:
print("Error: No images found in " + dataDir)
exit(0)
for dir in random.sample(dirs, n_labels_tot):
category = os.path.basename(dir)[0:1]
name = os.path.basename(dir)[2:].replace("_", " ")
labels.append(name)
categories.append(category)
for file in random.sample(glob.glob(dir + '/*.jpg'), rounds):
imageFiles.append(file)
names = random.sample(allNames, n_labels_teach)
models = [BaselineModel(),
CollieModelWOScaling(),
CollieFullModel(),
FewShotLearner()]
for round in range(rounds):
print(f" Round {round}")
imageFilesRound = []
for i in range(n_labels_tot):
imageFilesRound.append(imageFiles[i * rounds + round])
imgEmbedsRound = collie.encode_images(imageFilesRound)
for i in range(n_labels_teach):
# Now we evaluate performance on the teached objects using new names
phrase = names[i]
for model in models:
mrr = collie.mrr(model.find_best(phrase, imgEmbedsRound), i)
if mrr is not None:
out.write(f"{iteration}\t{round}\t{phrase}\t{model.name}\t{categories[i]}\t{labels[i]}\tcl\t{mrr}\n")
# Now we evaluate performance on the teached objects using original names
phrase = labels[i]
for model in models:
mrr = collie.mrr(model.find_best(phrase, imgEmbedsRound), i)
if mrr is not None:
out.write(f"{iteration}\t{round}\t{phrase}\t{model.name}\t{categories[i]}\t{labels[i]}\tzs-same\t{mrr}\n")
# Now we evaluate performance on not-teached objects using original names
for i in range(n_labels_teach, n_labels_teach + n_labels_test):
phrase = labels[i]
for model in models:
mrr = collie.mrr(model.find_best(phrase, imgEmbedsRound), i)
if mrr is not None:
out.write(f"{iteration}\t{round}\t{phrase}\t{model.name}\t{categories[i]}\t{labels[i]}\tzs-other\t{mrr}\n")
# Now we teach it the new names
for model in models:
for i in range(n_labels_teach):
phrase = names[i]
model.teach(phrase, imgEmbedsRound, i)
model.train()
del imgEmbedsRound
out.close()