-
Notifications
You must be signed in to change notification settings - Fork 0
/
detect_duplicates.py
60 lines (47 loc) · 1.43 KB
/
detect_duplicates.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
import os, shutil
from pathlib import Path
from pprint import pprint
import click
import cv2
import numpy as np
import tensorflow as tf
import imagehash
from PIL import Image
import utils
from tqdm import tqdm
def detect_duplicates(experiment_dir: str):
source_dir = Path(experiment_dir)
data_dir = source_dir / "data"
image_paths = list((data_dir).glob("**/*.png"))
print(f"Num: {len(image_paths)}")
hashes = dict()
hash_funcs = [
('ahash', imagehash.average_hash),
('phash', imagehash.phash),
('dhash', imagehash.dhash),
('whash-haar', imagehash.whash),
('whash-db4', lambda img: imagehash.whash(img, mode='db4')),
('cropresistant', imagehash.crop_resistant_hash),
]
for image_path in tqdm(image_paths):
with Image.open(str(image_path)) as img:
computed_hash = imagehash.average_hash(img, 32)
if hash in hashes:
hashes[computed_hash].append(image_path)
print(f"Duplicate found")
else:
hashes[computed_hash] = [image_path]
@click.command(help="Detect duplicates.")
@click.option(
"-e",
"--experiment-dir",
type=click.STRING,
required=True,
help="Path to experiment directory."
)
def main(experiment_dir: str):
detect_duplicates(
experiment_dir=experiment_dir,
)
if __name__ == "__main__":
main() # pylint: disable=no-value-for-parameter