import os
import shutil

# Define the paths
image_directory = './data/coco_det/images/semantics'  # The directory containing the images
txt_file_path = './list_union3.txt'       # The text file with indexes
new_directory = './data/coco_det/images/union3'      # The directory to copy matching images

# Create the new directory if it doesn't exist
if not os.path.exists(new_directory):
    os.makedirs(new_directory)

# Load the comma-separated indexes from the text file
with open(txt_file_path, 'r') as file:
    content = file.read().strip()  # Read the entire file content
    indexes = content.split(', ')   # Split by commas to get a list of indexes

# Iterate over the files in the image directory
for image_file in os.listdir(image_directory):
    if image_file.endswith(".jpg"):
        # Parse the appid and index from the image filename (example: 269170_1.jpg)
        image_appid, image_index = image_file.split("_")[0], image_file.split("_")[1].split(".")[0].zfill(3)
        
        # Combine appid and index to match the format in the text file (example: 269170001)
        image_key = f"{image_appid}{image_index}"
        
        # If the constructed key matches any entry in the text file, copy the image
        if image_key in indexes:
            source_path = os.path.join(image_directory, image_file)
            destination_path = os.path.join(new_directory, image_file)
            
            # Copy the image
            shutil.copy(source_path, destination_path)
            print(f"Copied: {image_file}")

print("All matching images have been copied.")