vcs-python-repository/YOLO python/clean-predictions.py
2024-04-13 17:16:17 -07:00

49 lines
No EOL
1.7 KiB
Python

import os
import shutil
import sys
def find_matching_images(text_folder, image_folder):
# Get list of text files in the text folder
text_files = [f for f in os.listdir(text_folder) if f.endswith('.txt')]
for text_file in text_files:
# Extract the number X from the text file name
file_number = text_file.split('_')[1].split('.')[0]
# Read the text file
with open(os.path.join(text_folder, text_file), 'r') as file:
lines = file.readlines()
# Filter lines with heights less than 0.01
filtered_lines = []
for line in lines:
parts = line.strip().split(' ')
if len(parts) == 5 and float(parts[4]) >= 0.01:
filtered_lines.append(line)
# Check if there are any rows left
if len(filtered_lines) == 0:
# Construct the image file name
image_file = file_number + '.jpg'
# Check if the image file exists in the image folder
if image_file in os.listdir(image_folder):
# If the image exists, print a message and delete it
print(f"Deleting image and text file for {text_file}")
os.remove(os.path.join(image_folder, image_file))
os.remove(os.path.join(text_folder, text_file))
else:
# Rewrite the text file with filtered lines
with open(os.path.join(text_folder, text_file), 'w') as file:
file.writelines(filtered_lines)
if __name__ == "__main__":
if len(sys.argv) != 3:
print("Usage: python script_name.py text_folder image_folder")
sys.exit(1)
text_folder = sys.argv[1]
image_folder = sys.argv[2]
find_matching_images(text_folder, image_folder)