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)