131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
import os
|
|
import struct
|
|
from tqdm import tqdm
|
|
from collections import Counter
|
|
import hashlib
|
|
|
|
def get_top_identities_binary(rec_path, idx_path, top_n=51):
|
|
"""
|
|
Pass 1: Scans the actual BINARY HEADERS in the .rec file.
|
|
This is the only way to be 100% sure which image belongs to whom.
|
|
"""
|
|
identity_counts = Counter()
|
|
|
|
with open(idx_path, 'r') as f:
|
|
offsets = [int(line.strip().split('\t')[1]) for line in f.readlines()]
|
|
|
|
print("Pass 1: Scanning binary headers to count identities...")
|
|
with open(rec_path, 'rb') as f:
|
|
for offset in tqdm(offsets):
|
|
f.seek(offset)
|
|
header_bin = f.read(32) # Read enough for the header
|
|
if len(header_bin) < 32: continue
|
|
|
|
# MXNet Header format: [Flag, Label (float), ID, ID]
|
|
# The label is at offset 12 (float32)
|
|
label = int(struct.unpack('f', header_bin[12:16])[0])
|
|
identity_counts[label] += 1
|
|
|
|
top_stats = identity_counts.most_common(top_n)
|
|
top_labels = {label for label, count in top_stats}
|
|
|
|
print(f"\nTop {top_n} Identities by Binary Label:")
|
|
for label, count in top_stats:
|
|
print(f"ID: {label:<10} | Count: {count:<10}")
|
|
|
|
return top_labels
|
|
|
|
def extract_selected_binary(rec_path, idx_path, output_dir, top_labels):
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
|
|
with open(idx_path, 'r') as f:
|
|
offsets = [int(line.strip().split('\t')[1]) for line in f.readlines()]
|
|
|
|
print(f"\nPass 2: Extracting verified images...")
|
|
|
|
# NEW: Keep track of how many images we've saved for each ID
|
|
# to avoid overwriting files.
|
|
save_counters = {label: 0 for label in top_labels}
|
|
total_extracted = 0
|
|
|
|
with open(rec_path, 'rb') as f:
|
|
for offset in tqdm(offsets):
|
|
f.seek(offset)
|
|
header_bin = f.read(32)
|
|
if len(header_bin) < 32: continue
|
|
|
|
label = int(struct.unpack('f', header_bin[12:16])[0])
|
|
|
|
if label not in top_labels:
|
|
continue
|
|
|
|
# Read image content
|
|
_, length_flag = struct.unpack('II', header_bin[:8])
|
|
content_length = length_flag & ((1 << 31) - 1)
|
|
content = f.read(content_length)
|
|
|
|
img_start = content.find(b'\xff\xd8')
|
|
if img_start == -1: continue
|
|
|
|
target_folder = os.path.join(output_dir, str(label))
|
|
os.makedirs(target_folder, exist_ok=True)
|
|
|
|
# Use the counter for this specific label
|
|
current_count = save_counters[label]
|
|
img_filename = f"{current_count}.jpg"
|
|
img_path = os.path.join(target_folder, img_filename)
|
|
if(current_count > 405):
|
|
continue
|
|
|
|
with open(img_path, 'wb') as img_f:
|
|
img_f.write(content[img_start:])
|
|
|
|
save_counters[label] += 1
|
|
total_extracted += 1
|
|
|
|
print(f"\nDone! Extracted {total_extracted} total images.")
|
|
|
|
|
|
|
|
def remove_duplicates(root_dir):
|
|
hashes = {} # hash -> first_filepath
|
|
duplicates_removed = 0
|
|
|
|
# Walk through every identity folder
|
|
for subdir, dirs, files in os.walk(root_dir):
|
|
for filename in tqdm(files, desc=f"Checking {os.path.basename(subdir)}"):
|
|
filepath = os.path.join(subdir, filename)
|
|
|
|
# Calculate MD5 hash of the file
|
|
with open(filepath, 'rb') as f:
|
|
file_hash = hashlib.md5(f.read()).hexdigest()
|
|
|
|
if file_hash in hashes:
|
|
# We've seen this image before!
|
|
os.remove(filepath)
|
|
duplicates_removed += 1
|
|
else:
|
|
hashes[file_hash] = filepath
|
|
|
|
print(f"\nClean-up complete. Removed {duplicates_removed} duplicate images.")
|
|
|
|
|
|
'''
|
|
if __name__ == "__main__":
|
|
# Point this to your unpacked Top 50 folder
|
|
target_dir = "./datasets/casia-set"
|
|
remove_duplicates(target_dir)
|
|
'''
|
|
if __name__ == "__main__":
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
REC = os.path.join(base_dir, '../data/casia-set', 'train.rec')
|
|
IDX = os.path.join(base_dir, '../data/casia-set', 'train.idx')
|
|
OUT = os.path.join(base_dir, '../data/casia-set')
|
|
|
|
# Step 1: Trust the binary, not the text file
|
|
top_verified_labels = get_top_identities_binary(REC, IDX, top_n=50)
|
|
|
|
# Step 2: Extract
|
|
extract_selected_binary(REC, IDX, OUT, top_verified_labels)
|
|
|