1
from typing import Optional
2
import cv2
3
import shutil
4
from pathlib import Path
5
from skimage.metrics import structural_similarity as ssim
6
from rich.progress import Progress
7
from rich.console import Console
8
from annoy import AnnoyIndex
9
import numpy as np
10

11
RAW_DIR = Path("data/step-1-deduplication")
12
DUP_DIR = RAW_DIR / "duplicate"
13
INV_DIR = RAW_DIR / "invalid"
14

15
DUP_DIR.mkdir(exist_ok=True)
16
INV_DIR.mkdir(exist_ok=True)
17

18
IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff", ".webp"}
19
SIMILARITY_THRESHOLD = 0.95  # 精确判断阈值（SSIM）
20
VECTOR_SIZE = 128  # ORB向量长度
21
NEAREST_NEIGHBORS = 5  # 查询最近邻数量
22

23
console = Console()
24

25

26
def is_valid_image(path: Path) -> bool:
27
    try:
28
        img = cv2.imread(str(path))
29
        return img is not None
30
    except Exception:
31
        return False
32

33

34
def load_gray_image(path: Path, size=256):
35
    img = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
36
    if img is not None:
37
        img = cv2.resize(img, (size, size))
38
    return img
39

40

41
def compare_images(img1, img2):
42
    return ssim(img1, img2)
43

44

45
def move_with_txt(src: Path, dst_dir: Path, keep: Optional[Path] = None):
46
    dst_dir.mkdir(exist_ok=True)
47
    dst_path = dst_dir / src.name
48
    shutil.move(str(src), dst_path)
49

50
    src_txt = src.with_suffix(".txt")
51
    dst_txt = dst_path.with_suffix(".txt")
52

53
    if src_txt.exists():
54
        shutil.move(str(src_txt), dst_txt)
55
        # 如果有保留的文件但它没有标签文件，则复制一份
56
        if keep is not None:
57
            keep_txt = keep.with_suffix(".txt")
58
            if not keep_txt.exists():
59
                shutil.copy(dst_txt, keep_txt)
60
                console.log(f"[TXT-COPIED] {src_txt} → {keep_txt}")
61

62

63
def get_orb_vector(path: Path):
64
    img = load_gray_image(path)
65
    if img is None:
66
        return np.zeros(VECTOR_SIZE, dtype=np.float32)
67
    orb = cv2.ORB_create()
68
    kp, des = orb.detectAndCompute(img, None)
69
    if des is None:
70
        return np.zeros(VECTOR_SIZE, dtype=np.float32)
71
    vec = des.flatten()
72
    if len(vec) < VECTOR_SIZE:
73
        vec = np.pad(vec, (0, VECTOR_SIZE - len(vec)))
74
    else:
75
        vec = vec[:VECTOR_SIZE]
76
    return vec.astype(np.float32)
77

78

79
def main():
80
    files = [
81
        f for f in RAW_DIR.iterdir() if f.is_file() and f.suffix.lower() in IMG_EXTS
82
    ]
83
    vectors = []
84
    valid_files: list[Path] = []
85

86
    # 1. 构建特征向量
87
    with Progress(console=console) as progress:
88
        task = progress.add_task("[cyan]Extracting features...", total=len(files))
89
        for file in files:
90
            progress.update(task, description=f"[cyan]Processing {file.name}")
91
            if not is_valid_image(file):
92
                move_with_txt(file, INV_DIR)
93
                console.log(f"[INVALID] {file}")
94
                progress.advance(task)
95
                continue
96
            vec = get_orb_vector(file)
97
            vectors.append(vec)
98
            valid_files.append(file)
99
            progress.advance(task)
100

101
    if not valid_files:
102
        console.log("No valid images found.")
103
        return
104

105
    # 2. 构建 Annoy 索引
106
    dim = VECTOR_SIZE
107
    index = AnnoyIndex(dim, "euclidean")
108
    for i, vec in enumerate(vectors):
109
        index.add_item(i, vec)
110
    index.build(10)  # 可调
111
    processed_indices = set()
112
    duplicate_map = {}
113
    console.log(f"Annoy index built with {len(valid_files)} images.")
114

115
    with Progress(console=console) as progress:
116
        task = progress.add_task("[green]Finding duplicates...", total=len(valid_files))
117
        for i, file in enumerate(valid_files):
118
            if file in processed_indices or not file.exists():
119
                progress.advance(task)
120
                continue
121
            progress.update(task, description=f"[green]Checking {file.name}")
122
            if i in processed_indices:
123
                progress.advance(task)
124
                continue
125

126
            # 查询最近邻
127
            nearest = index.get_nns_by_item(i, NEAREST_NEIGHBORS)
128
            found_duplicate = False
129
            img_gray = load_gray_image(file)
130
            for j in nearest:
131
                if j == i or j in processed_indices:
132
                    continue
133
                other_file = valid_files[j]
134
                if not other_file.exists():
135
                    processed_indices.add(j)
136
                    continue
137
                other_gray = load_gray_image(other_file)
138
                score = compare_images(img_gray, other_gray)
139
                if score >= SIMILARITY_THRESHOLD:
140
                    # 保留体积更大的
141
                    if file.stat().st_size >= other_file.stat().st_size:
142
                        move_with_txt(other_file, DUP_DIR, keep=file)
143
                        console.log(
144
                            f"[DUPLICATE] {other_file} → {DUP_DIR}, kept {file} (SSIM={score:.3f})"
145
                        )
146
                        duplicate_map[DUP_DIR / other_file.name] = file
147
                        processed_indices.add(j)
148
                    else:
149
                        move_with_txt(file, DUP_DIR, keep=other_file)
150
                        console.log(
151
                            f"[DUPLICATE] {file} → {DUP_DIR}, kept {other_file} (SSIM={score:.3f})"
152
                        )
153
                        found_duplicate = True
154
                        break
155
            if not found_duplicate:
156
                processed_indices.add(i)
157
            progress.advance(task)
158
    for k, v in duplicate_map.items():
159
        console.log(f"{k} keep: {v}")
160
    print(duplicate_map)
161

162

163
if __name__ == "__main__":
164
    main()