migration and splitting AareLC
This commit is contained in:
@@ -0,0 +1,188 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
import glob
|
||||
import yaml
|
||||
from pathlib import Path
|
||||
|
||||
def resize_to_screen(img, screen_width=1920, screen_height=1080):
|
||||
"""Resize the image to fit the screen while maintaining the aspect ratio."""
|
||||
height, width, _ = img.shape
|
||||
scaling_factor = min(screen_width / width, screen_height / height)
|
||||
new_width = int(width * scaling_factor)
|
||||
new_height = int(height * scaling_factor)
|
||||
resized_img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
|
||||
return resized_img, scaling_factor
|
||||
|
||||
def draw_and_record_boxes(image_path, screen_width=1920, screen_height=1080):
|
||||
"""Draw four boxes and record normalized YOLO box data."""
|
||||
img = cv2.imread(image_path)
|
||||
|
||||
# Resize image to fit the screen
|
||||
resized_img, scale = resize_to_screen(img, screen_width/2, screen_height/2)
|
||||
height, width, _ = img.shape # Original dimensions (for normalization)
|
||||
|
||||
while True:
|
||||
|
||||
boxes_data = [] # List to store box info
|
||||
preview_img = resized_img.copy()
|
||||
|
||||
# Helper for selecting and recording a box
|
||||
def select_and_record_box(box_name, color, class_id):
|
||||
box = cv2.selectROI(box_name, resized_img, fromCenter=False, showCrosshair=True)
|
||||
if box == (0, 0, 0, 0): # no selection
|
||||
return None
|
||||
x, y, w, h = [int(coord / scale) for coord in box] # Scale back to original size
|
||||
rx, ry, rw, rh = [int(coord) for coord in box] # use resized coords directly
|
||||
cv2.rectangle(preview_img, (rx, ry), (rx + rw, ry + rh), color, 2)
|
||||
#cv2.rectangle(resized_img, (x, y), (x + w, y + h), color, 2)
|
||||
center_x = (x + w / 2) / width
|
||||
center_y = (y + h / 2) / height
|
||||
norm_w = w / width
|
||||
norm_h = h / height
|
||||
return [class_id, center_x, center_y, norm_w, norm_h]
|
||||
|
||||
# Boxes
|
||||
|
||||
box = select_and_record_box(f"Select Class 0: {class_dict[0].upper()} GREEN Box", (0, 255, 0), 0)
|
||||
if box: boxes_data.append(box)
|
||||
|
||||
box = select_and_record_box(f"Select Class 1: {class_dict[1].upper()} RED Box", (0, 0, 255), 1)
|
||||
if box: boxes_data.append(box)
|
||||
|
||||
box = select_and_record_box(f"Select Class 2: {class_dict[2].upper()} BLUE Box", (255, 0, 0), 2)
|
||||
if box: boxes_data.append(box)
|
||||
|
||||
box = select_and_record_box(f"Select Class 3: {class_dict[3].upper()} YELLOW Box", (0, 255, 255), 3)
|
||||
if box: boxes_data.append(box)
|
||||
|
||||
# Show annotated image
|
||||
cv2.imshow("Annotated Image (Press: [a]=accept, [r]=restart, [q/ESC]=quit)", preview_img)
|
||||
key = cv2.waitKey(0) & 0xFF # Mask to 8-bit
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
if key in [ord('a'), 13]: # 'a' or ENTER = accept
|
||||
return boxes_data, False, preview_img
|
||||
elif key == ord('r'): # restart selection
|
||||
print("Restarting annotation for this image...")
|
||||
continue
|
||||
elif key in [27, ord('q')]: # ESC or 'q' = quit all
|
||||
return boxes_data, True, preview_img
|
||||
elif key == ord('s'): # s = skip current image
|
||||
print("Skipping this image...")
|
||||
return [], False, preview_img
|
||||
else:
|
||||
print("Unrecognized key. Press [a]=accept, [r]=restart, [q]=quit.")
|
||||
|
||||
return boxes_data, False, preview_img
|
||||
|
||||
|
||||
def process_images_for_yolo(input_folder, labels_folder, screen_width=1920, screen_height=1080, overwrite_existing=None):
|
||||
"""Process images for YOLO by saving labeled bounding boxes for each image."""
|
||||
# Supported image extensions
|
||||
image_extensions = ('*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tif', '*.tiff')
|
||||
|
||||
# Collect all image file paths from the input folder
|
||||
image_paths = []
|
||||
for ext in image_extensions:
|
||||
image_paths.extend(glob.glob(os.path.join(input_folder, ext)))
|
||||
|
||||
if not image_paths:
|
||||
print(f"No images found in the input folder: {input_folder}")
|
||||
return
|
||||
|
||||
# Create the labels folder if it doesn't exist
|
||||
if not os.path.exists(labels_folder):
|
||||
os.makedirs(labels_folder)
|
||||
|
||||
print("\n--- YOLO Image Annotation Controls ---")
|
||||
print("[ESC] or [q] → quit annotation completely")
|
||||
print("[c] - cancel process")
|
||||
print("[r] → restart current image")
|
||||
print("[s] → skip current image")
|
||||
print("[any other] → save boxes & move to next image")
|
||||
print("---------------------------------------\n")
|
||||
|
||||
# Iterate over each image
|
||||
for image_path in image_paths:
|
||||
print(f"Processing: {image_path}")
|
||||
image_name = os.path.splitext(os.path.basename(image_path))[0]
|
||||
output_file_path = os.path.join(labels_folder, f"{image_name}.txt")
|
||||
|
||||
if os.path.exists(output_file_path):
|
||||
if overwrite_existing is True:
|
||||
print(f"overwriteing exisiting label for {image_name}.txt")
|
||||
elif overwrite_existing is False:
|
||||
print(f"Skipping {image_name} as already labeled")
|
||||
continue
|
||||
else:
|
||||
print(f"⚠️ Label already exists for {image_name}.txt")
|
||||
choice = input("[o] → overwrite, [s] → skip, [q/ESC] → quit: "). strip().lower()
|
||||
|
||||
if choice == "s":
|
||||
print(f"Skipping {image_name} (already labeled).")
|
||||
continue
|
||||
elif choice == "q":
|
||||
print("Exiting early...")
|
||||
return
|
||||
elif choice == "o":
|
||||
print(f"Overwriting label for {image_name}...")
|
||||
else:
|
||||
print("Unknown choice, skipping this image.")
|
||||
continue
|
||||
|
||||
print(f"\nProcessing: {image_path}")
|
||||
|
||||
while True:
|
||||
# Get normalized box data from user interaction
|
||||
boxes_data, exit_flag, preview_img = draw_and_record_boxes(image_path, screen_width, screen_height)
|
||||
|
||||
if exit_flag: # Exit the annotation process if Ctrl+W is pressed
|
||||
print("Exiting image annotation early...")
|
||||
return
|
||||
|
||||
if boxes_data: # Write the annotations for this image file if boxes are drawn
|
||||
image_name = os.path.splitext(os.path.basename(image_path))[0] # File name without extension
|
||||
output_file_path = os.path.join(labels_folder, f"{image_name}.txt")
|
||||
|
||||
# Write each box to its corresponding text file in YOLO format
|
||||
with open(output_file_path, 'w') as f:
|
||||
for box in boxes_data:
|
||||
class_id, center_x, center_y, norm_w, norm_h = box
|
||||
f.write(f"{class_id} {center_x:.6f} {center_y:.6f} {norm_w:.6f} {norm_h:.6f}\n")
|
||||
|
||||
#cv2.imshow(f"Saved Annotation: {image_name}", preview_img)
|
||||
#cv2.waitKey(1500) # wait 1.5 seconds or until key is pressed
|
||||
#cv2.destroyAllWindows()
|
||||
|
||||
break
|
||||
|
||||
else:
|
||||
print("skipping this image...")
|
||||
break
|
||||
|
||||
print(f"Annotations saved to {labels_folder}")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Define input folder containing images
|
||||
input_folder = r"/Users/gotthardg/Volumes/ra_work/Sept/annotations" #"/Users/duan_j/Applications/alc/tests/yolo/nodetec/pin2" # Change to your folder path
|
||||
#overwrite flag to say whether to overwrite existing labels, skip all labeled or enable user choice.
|
||||
overwrite_flag = True
|
||||
# Define the labels output folder for YOLO format
|
||||
labels_folder = f"{input_folder}labels" # Change to your desired folder
|
||||
|
||||
# Path to your yaml file
|
||||
yaml_file = str(Path(__file__).resolve().parents[2] / "config" / "dataset.yaml")
|
||||
|
||||
with open(yaml_file, "r") as f:
|
||||
data = yaml.safe_load(f)
|
||||
|
||||
# Extract names dictionary
|
||||
class_dict = data.get("names", {})
|
||||
|
||||
print(class_dict[0])
|
||||
|
||||
# Process images and save YOLO-compatible annotations
|
||||
process_images_for_yolo(input_folder, labels_folder, overwrite_existing=overwrite_flag)
|
||||
Reference in New Issue
Block a user