189 lines
7.8 KiB
Python
189 lines
7.8 KiB
Python
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)
|