import cv2
import sys
import os

def get_yolo_coordinates(image_path):
    """
    Opens an image, lets the user draw a box, and prints the YOLO format coordinates.
    """
    if not os.path.exists(image_path):
        print(f"Error: File {image_path} not found.")
        return

    # Load image
    img = cv2.imread(image_path)
    if img is None:
        print("Error: Could not read image.")
        return

    height, width, _ = img.shape
    
    print("---------------------------------------------------------")
    print(f"Loaded {image_path} ({width}x{height})")
    print("INSTRUCTIONS:")
    print("1. Draw a box around the object using your mouse.")
    print("2. Press ENTER or SPACE to confirm the box.")
    print("3. Press 'c' to cancel.")
    print("---------------------------------------------------------")

    # Select ROI
    # fromCenter=False, showCrosshair=True
    r = cv2.selectROI("Draw Box (Press Enter to Confirm)", img, fromCenter=False, showCrosshair=True)
    cv2.destroyAllWindows()
    
    # r is (x, y, w, h) in pixels
    x_pixel, y_pixel, w_pixel, h_pixel = r

    if w_pixel == 0 or h_pixel == 0:
        print("No box selected.")
        return

    # Convert to YOLO format (Normalized)
    # center_x, center_y, width, height
    
    center_x = (x_pixel + (w_pixel / 2)) / width
    center_y = (y_pixel + (h_pixel / 2)) / height
    norm_w = w_pixel / width
    norm_h = h_pixel / height

    # Limit precision to 6 decimal places
    print(f"\nSUCCESS! Here is your YOLO label line:")
    print(f"---------------------------------------------------------")
    print(f"<class_id> {center_x:.6f} {center_y:.6f} {norm_w:.6f} {norm_h:.6f}")
    print(f"---------------------------------------------------------")
    print("Replace <class_id> with:")
    print("0 -> if it is a Traffic Sign")
    print("1 -> if it is a Pothole")
    print("2 -> if it is a Manhole")
    print(f"---------------------------------------------------------")

if __name__ == "__main__":
    if len(sys.argv) > 1:
        path = sys.argv[1]
    else:
        path = input("Enter the path to your image: ").strip('"')
    
    get_yolo_coordinates(path)