提高训练样本质量，去除相似图片

提高训练样本质量，去除相似图片

使用了fiftyone，数据集使用yolo，多个数据集挑选质量最高的一批图片，然后合并，剩下的图片作为测试

import fiftyone as fo
import fiftyone.zoo as foz
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import os
import shutil
from pathlib import Path

def fiftyone_image_deduplication_yolo(in_dir,out_dir,limit=50,thresh = 0.92):    
    path_img = os.path.join(in_dir, "images")
    path_label = os.path.join(in_dir, "labels")
    path_label_class = os.path.join(path_label, "classes.txt")
    
    path_out_img = os.path.join(out_dir, "images")
    path_out_label = os.path.join(out_dir, "labels")
    path_out_label_class = os.path.join(path_out_label, "classes.txt")
    path_out_test = os.path.join(out_dir, "test")
    
    if os.path.exists(out_dir):        
        shutil.rmtree(out_dir)
    os.makedirs(path_out_img)
    os.makedirs(path_out_label)
    os.makedirs(path_out_test)
    
    dataset = fo.Dataset.from_images_dir(path_img)
    
    #dataset = fo.Dataset.from_dir(
    #    dataset_dir=in_dir,
    #    dataset_type=fo.types.ImageDirectory,
    #    name="my-image-deduplication",
    #)
    
    model = foz.load_zoo_model("mobilenet-v2-imagenet-torch")
    embeddings = dataset.compute_embeddings(model)
    #print(embeddings.shape)

    similarity_matrix = cosine_similarity(embeddings)
    #print(similarity_matrix.shape)
    #print(similarity_matrix)

    n = len(similarity_matrix)
    similarity_matrix = similarity_matrix - np.identity(n)
    #print(similarity_matrix)

    id_map = [s.id for s in dataset.select_fields(["id"])]
    for idx, sample in enumerate(dataset):
        sample["max_similarity"] = similarity_matrix[idx].max()
        sample.save()

    samples_to_remove = set()
    samples_to_keep = set()
    for idx, sample in enumerate(dataset):
        if sample.id not in samples_to_remove:
            # Keep the first instance of two duplicates
            samples_to_keep.add(sample.id)

            dup_idxs = np.where(similarity_matrix[idx] > thresh)[0]
            for dup in dup_idxs:
                # We kept the first instance so remove all other duplicates
                samples_to_remove.add(id_map[dup])

            if len(dup_idxs) > 0:
                sample.tags.append("has_duplicates")
                sample.save()

        else:
            sample.tags.append("duplicate")
            sample.save()

    print("samples total:",len(dataset))
    print("remove:",len(samples_to_remove))
    print("keep:",len(samples_to_keep))
    
    dataset.delete_samples(list(samples_to_remove))
    #print(len(dataset))

    view = dataset.sort_by("max_similarity", reverse=False).limit(limit)
    select_img = []
    for idx, sample in enumerate(view):
        print(idx,sample.filepath,sample.max_similarity)
        bname = os.path.basename(sample.filepath)
        select_img.append(bname)
        sptext = os.path.splitext(bname)
        org_file_img = sample.filepath
        org_file_label = os.path.join(path_label, sptext[0]+'.txt')
        out_file_img = os.path.join(path_out_img, bname)
        out_file_label = os.path.join(path_out_label, sptext[0]+'.txt')
        shutil.copy(org_file_img, out_file_img)
        shutil.copy(org_file_label, out_file_label)
    
    #挑选剩下的作为测试
    all_img = []
    for file in os.listdir(path_img):        
        all_img.append(file)

    test_img = list(set(all_img)-set(select_img))
    print("test img:",len(test_img))
    for file in test_img:
        file_img = os.path.join(path_img, file)
        shutil.copy(file_img, path_out_test)
    
def dir_cut(src,dst,rm=True):
    if os.path.exists(dst) and rm:   
        shutil.rmtree(dst)
    if not os.path.exists(dst):
        os.makedirs(dst)
    for file in os.listdir(src):
        file_src = os.path.join(src, file)
        shutil.copy(file_src, dst)
    shutil.rmtree(src)
    
def deduplication_copy(src,dst):
    path_img = os.path.join(src, "images")
    path_label = os.path.join(src, "labels")
    path_test = os.path.join(src, "test")
    path_out_img = os.path.join(dst, "images")
    path_out_label = os.path.join(dst, "labels")
    path_out_test = os.path.join(dst, "test")
    dir_cut(path_img, path_out_img, rm=False)
    dir_cut(path_label, path_out_label, rm=False)
    dir_cut(path_test, path_out_test, rm=False)

def merge_deduplication():
    # 输入输出是 yolo 格式，images和labels目录
    fiftyone_image_deduplication_yolo(
        in_dir="/workspace/data/small_flame_fog_smoke/org/smoke",
        out_dir="/workspace/data/small_flame_fog_smoke/tmp/smoke",
        limit=300)

    fiftyone_image_deduplication_yolo(
        in_dir="/workspace/data/small_flame_fog_smoke/org/fog",
        out_dir="/workspace/data/small_flame_fog_smoke/tmp/fog",
        limit=50)

    fiftyone_image_deduplication_yolo(
        in_dir="/workspace/data/small_flame_fog_smoke/org/small_flame",
        out_dir="/workspace/data/small_flame_fog_smoke/tmp/small_flame",
        limit=70)

    deduplication_dir = Path("/workspace/data/small_flame_fog_smoke/deduplication/yolo")
    if os.path.exists(deduplication_dir):        
        shutil.rmtree(deduplication_dir)
    os.makedirs(deduplication_dir)
    deduplication_copy(Path("/workspace/data/small_flame_fog_smoke/tmp/smoke"),deduplication_dir)
    deduplication_copy(Path("/workspace/data/small_flame_fog_smoke/tmp/fog"),deduplication_dir)
    deduplication_copy(Path("/workspace/data/small_flame_fog_smoke/tmp/small_flame"),deduplication_dir)
    shutil.rmtree(Path("/workspace/data/small_flame_fog_smoke/tmp"))
    shutil.copy(Path("/workspace/data/small_flame_fog_smoke/org/classes.txt"), os.path.join(deduplication_dir, "labels", "classes.txt"))

    path_test = os.path.join(deduplication_dir, "test")
    path_out_test = Path("/workspace/data/small_flame_fog_smoke/test")
    dir_cut(path_test,path_out_test)

merge_deduplication()
print("done")