Spaces:
Sleeping
Sleeping
File size: 1,172 Bytes
8424a77 |
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 |
# takes a pickle, takes a file with md5sums, outputs the pickle without hash dupes.
import sys
import pickle
import numpy as np
pickle_filename, md5_filename, output_pickle_filename = sys.argv[1:]
hashes = {}
for line in open(md5_filename, "r"):
line = line.strip("\n")
md5 = line.split()[0]
assert len(md5) == 32
assert line[32:34] == " "
filename = line[34:]
hashes[filename] = md5
print(len(hashes), "hashes read")
data = pickle.load(open(pickle_filename, "rb"))
print(len(data["embeddings"]), "embeddings read")
filenames = data["filenames"]
collected_indices = []
collected_md5s = set()
for i in range(len(filenames)):
filename = filenames[i]
md5 = hashes[filename] # not the hash of the filename, the hash of the file content.
if md5 not in collected_md5s:
collected_indices.append(i)
collected_md5s.add(md5)
print(len(collected_indices), "unique hashes")
filenames = np.array(filenames)
data["filenames"] = filenames[collected_indices].tolist()
data["embeddings"] = data["embeddings"][collected_indices]
assert "thumbs" not in data
with open(output_pickle_filename, "wb") as f:
pickle.dump(data, f)
|