You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
112 lines
5.7 KiB
112 lines
5.7 KiB
from ultralytics import YOLO
|
|
import os
|
|
import shutil
|
|
from PIL import Image
|
|
|
|
# --- Configuration ---
|
|
# Load your trained model
|
|
model = YOLO('./best.pt') # Adjust 'train8' if needed
|
|
|
|
# Directory with new, unannotated images (your input)
|
|
unlabeled_image_dir = './untrained_images'
|
|
|
|
# Base directory for YOLOv8's output. YOLOv8 will create 'predictX' folders inside this.
|
|
# Example: /TempReview/predict, /TempReview/predict1, etc.
|
|
yolov8_project_dir = './.yolo_yemp' # This is where your 'predictX' folders are generated
|
|
|
|
# The final, flat directory where images and labels will be moved for LabelImg
|
|
flat_output_dir = './for_labelimg_review'
|
|
|
|
# --- Ensure input directory exists ---
|
|
os.makedirs(unlabeled_image_dir, exist_ok=True) # Make sure this exists if you haven't uploaded images yet
|
|
|
|
# Get all image files from the unlabeled directory
|
|
image_files = [f for f in os.listdir(unlabeled_image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
|
|
|
|
if not image_files:
|
|
print(f"No image files found in '{unlabeled_image_dir}'. Please upload images there.")
|
|
else:
|
|
print(f"--- Step 1: Running YOLOv8 Inference on new images ---")
|
|
print(f"Running inference on {len(image_files)} new images...")
|
|
|
|
# Run inference on all images in the directory
|
|
# YOLOv8 will automatically create a new 'predictX' folder (e.g., predict, predict1, predict2)
|
|
# inside the yolov8_project_dir for each unique run.
|
|
results_list = model(unlabeled_image_dir,
|
|
#save=True, # Save images with plotted boxes
|
|
save_txt=True, # Save the YOLO .txt label files
|
|
conf=0.6, # Confidence threshold for predictions (adjust as needed)
|
|
project=yolov8_project_dir # This is your 'TempReview'
|
|
)
|
|
|
|
# Find the most recently created 'predictX' folder
|
|
# This assumes the latest run will be the one with the highest number
|
|
# or the one most recently modified.
|
|
predict_dirs = [d for d in os.listdir(yolov8_project_dir) if d.startswith('predict') and os.path.isdir(os.path.join(yolov8_project_dir, d))]
|
|
|
|
if not predict_dirs:
|
|
print(f"Error: No 'predictX' folders found in '{yolov8_project_dir}'. Inference might have failed.")
|
|
else:
|
|
# Sort by creation time (most recent first) or by name (highest number)
|
|
# Sorting by name (predict, predict1, predict10, predict2...) doesn't always work numerically.
|
|
# Sorting by modification time is safer.
|
|
latest_predict_dir_name = max(predict_dirs, key=lambda d: os.path.getmtime(os.path.join(yolov8_project_dir, d)))
|
|
|
|
yolov8_run_path = os.path.join(yolov8_project_dir, latest_predict_dir_name)
|
|
yolov8_images_path = yolov8_run_path
|
|
yolov8_labels_path = os.path.join(yolov8_run_path, 'labels')
|
|
|
|
print(f"YOLOv8 results saved to: '{yolov8_run_path}'")
|
|
|
|
print(f"\n--- Step 2: Flattening output structure for LabelImg ---")
|
|
os.makedirs(flat_output_dir, exist_ok=True)
|
|
|
|
# Move images
|
|
if os.path.exists(yolov8_images_path):
|
|
for img_file in os.listdir(yolov8_images_path):
|
|
# Only move files that are original image files (e.g., .jpg, .png)
|
|
if img_file.lower().endswith(('.jpg', '.jpeg', '.png')):
|
|
shutil.move(os.path.join(yolov8_images_path, img_file),
|
|
os.path.join(flat_output_dir, img_file))
|
|
print(f"Moved images to '{flat_output_dir}'")
|
|
else:
|
|
print(f"Warning: No images found at '{yolov8_images_path}'.")
|
|
|
|
# Process and move labels
|
|
if os.path.exists(yolov8_labels_path):
|
|
for label_file in os.listdir(yolov8_labels_path):
|
|
if label_file.lower().endswith('.txt'):
|
|
label_path_src = os.path.join(yolov8_labels_path, label_file)
|
|
|
|
# Read lines, parse class_id, and sort
|
|
with open(label_path_src, 'r') as f:
|
|
lines = f.readlines()
|
|
|
|
# Sort lines based on the first element (class_id) as an integer
|
|
# Handle potential errors if a line is malformed, though unlikely from YOLO.
|
|
try:
|
|
sorted_lines = sorted(lines, key=lambda line: int(line.strip().split(' ')[0]))
|
|
except ValueError as e:
|
|
print(f"Warning: Could not sort lines in {label_file} due to format error: {e}. Skipping sort for this file.")
|
|
sorted_lines = lines # Fallback to original order
|
|
|
|
# Write sorted lines to the destination file
|
|
label_path_dest = os.path.join(flat_output_dir, label_file)
|
|
with open(label_path_dest, 'w') as f:
|
|
f.writelines(sorted_lines)
|
|
|
|
# Remove the original label file after processing (optional, but keeps source clean)
|
|
os.remove(label_path_src)
|
|
print(f"Moved and sorted labels to '{flat_output_dir}'")
|
|
else:
|
|
print(f"Warning: No labels found at '{yolov8_labels_path}'.")
|
|
|
|
# Clean up the intermediate YOLOv8 output directory (optional, but recommended for Colab space)
|
|
# Be careful with shutil.rmtree - it deletes recursively!
|
|
if os.path.exists(yolov8_run_path):
|
|
print(f"Cleaning up temporary YOLOv8 output: '{yolov8_run_path}'")
|
|
shutil.rmtree(yolov8_run_path)
|
|
|
|
print(f"\n--- Process Complete! ---")
|
|
print(f"Your images and predicted .txt labels are now in a flat structure in: '{flat_output_dir}'")
|
|
print("You can now open LabelImg and point it to this directory to begin review and correction.")
|