非同步

Ding-yixia發表於2024-07-12
import cv2
import numpy as np
import os
import sqlite3
import concurrent.futures





def create_tables_and_insert_data(db_filename, high_similarity_pairs, low_similarity_pairs, unmatched_files_A, unmatched_files_B):
    conn = sqlite3.connect(db_filename)
    cursor = conn.cursor()
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS high_similarity_pairs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    path_A TEXT,
    path_B TEXT,
    similarity REAL
    )
    ''')

     
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS low_similarity_pairs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            path_A TEXT,
            path_B TEXT,
            similarity REAL
        )
    ''')

    cursor.execute('''
        CREATE TABLE IF NOT EXISTS unmatched_files_A (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            path TEXT
        )
    ''')

    cursor.execute('''
        CREATE TABLE IF NOT EXISTS unmatched_files_B (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            path TEXT
        )
    ''')

    for pair in high_similarity_pairs:
        cursor.execute('INSERT INTO high_similarity_pairs (path_A, path_B, similarity) VALUES (?, ?, ?)', pair)

    for pair in low_similarity_pairs:
        cursor.execute('INSERT INTO low_similarity_pairs (path_A, path_B, similarity) VALUES (?, ?, ?)', pair)

    for file in unmatched_files_A:
        cursor.execute('INSERT INTO unmatched_files_A (path) VALUES (?)', (file,))

    for file in unmatched_files_B:
        cursor.execute('INSERT INTO unmatched_files_B (path) VALUES (?)', (file,))

    conn.commit()
    conn.close()




def calculate_black_pixels_in_sectors(image, sector):
    _, image = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY_INV)
    height, width = image.shape
    center_x, center_y = width // 2, height // 2
    angle_step = 360 / sector
    black_pixel_counts = np.zeros(sector)

    for i in range(sector):
        start_angle = np.deg2rad(i * angle_step)
        end_angle = np.deg2rad((i + 1) * angle_step)
        y, x = np.ogrid[:height, :width]
        mask = ((x - center_x) * np.cos(start_angle) + (y - center_y) * np.sin(start_angle) >= 0) & \
               ((x - center_x) * np.cos(end_angle) + (y - center_y) * np.sin(end_angle) <= 0)
        black_pixel_counts[i] = np.sum(image[mask] == 0)

    return black_pixel_counts

def compare_images(image_path1, image_path2, pixel_threshold, sector_threshold, sector):
    image1 = cv2.imread(image_path1, cv2.IMREAD_GRAYSCALE)
    image2 = cv2.imread(image_path2, cv2.IMREAD_GRAYSCALE)
    black_pixels1 = calculate_black_pixels_in_sectors(image1, sector)
    black_pixels2 = calculate_black_pixels_in_sectors(image2, sector)
    diff = np.abs(black_pixels1 - black_pixels2)
    num_similar_sectors = np.sum(diff <= pixel_threshold)
    similarity_percentage = (num_similar_sectors / len(black_pixels1)) * 100
    is_within_sector_threshold = np.sum(diff > pixel_threshold) <= sector_threshold
    return similarity_percentage, is_within_sector_threshold

def compare_subdirectories(dirA, dirB, common_dir, pixel_threshold, sector_threshold, similarity_threshold, not_similarity_threshold, sector, db_dir):
    pathA = os.path.join(dirA, common_dir)
    pathB = os.path.join(dirB, common_dir)
    filesA = os.listdir(pathA)
    filesB = os.listdir(pathB)
    common_files = list(set(filesA).intersection(set(filesB)))
    results = {}
    high_similarity_pairs = []
    low_similarity_pairs = []
    for file_name in common_files:
        img_path_A = os.path.join(pathA, file_name)
        img_path_B = os.path.join(pathB, file_name)
        similarity, within_threshold = compare_images(img_path_A, img_path_B, pixel_threshold, sector_threshold, sector)
        results[file_name] = {'similarity': similarity, 'within_threshold': within_threshold}
        if similarity >= similarity_threshold:
            high_similarity_pairs.append((img_path_A, img_path_B, similarity))
        if similarity <= not_similarity_threshold:
            low_similarity_pairs.append((img_path_A, img_path_B, similarity))
    unmatched_files_A = list(set(filesA) - set(common_files))
    unmatched_files_B = list(set(filesB) - set(common_files))

    db_filename = os.path.join(db_dir, f"{common_dir}.db")
    create_tables_and_insert_data(
        db_filename,
        high_similarity_pairs,
        low_similarity_pairs,
        unmatched_files_A,
        unmatched_files_B
    )

    return {
        'results': results,
        'high_similarity_pairs': high_similarity_pairs,
        'low_similarity_pairs': low_similarity_pairs,
        'unmatched_files_A': unmatched_files_A,
        'unmatched_files_B': unmatched_files_B
    }

def compare_directories(dirA, dirB, pixel_threshold, sector_threshold, similarity_threshold, not_similarity_threshold, sector, db_dir):
    if not os.path.exists(dirA) or not os.path.exists(dirB):
        raise ValueError("One or both directories do not exist.")
    dirsA = os.listdir(dirA)
    dirsB = os.listdir(dirB)
    common_dirs = list(set(dirsA).intersection(set(dirsB)))
    results = {}

    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = {executor.submit(compare_subdirectories, dirA, dirB, common_dir, pixel_threshold, sector_threshold, similarity_threshold, not_similarity_threshold, sector, db_dir): common_dir for common_dir in common_dirs}
        for future in concurrent.futures.as_completed(futures):
            common_dir = futures[future]
            try:
                sub_results = future.result()
            except Exception as e:
                print(f'Error comparing directory {common_dir}: {e}')

# 假設 create_tables_and_insert_data 函式已實現





# 呼叫函式
dirA = r'imgA'
dirB = r'imgB'
pixel_threshold = 50
sector_threshold = 5
sector = 72
similarity_threshold = 70
not_similarity_threshold = 30
db_dir = r"Db_DIR"
if not os.path.exists(db_dir):
    os.makedirs(db_dir)

results = compare_directories(dirA, dirB, pixel_threshold, sector_threshold, similarity_threshold, not_similarity_threshold, sector, db_dir)


相關文章