interactive_selection.py

import cv2
import numpy as np
import json
from pathlib import Path

# 全局变量
ref_point = []
dragging = False
current_element = None
elements = []
element_types = ['Departure', 'Arrival', 'TrainNo', 'Class', 'Time', 'Price', 'Other']
current_type_index = 0

def draw_rectangle(img, pt1, pt2, color=(0, 255, 0), thickness=2):
    """Draw rectangle"""
    cv2.rectangle(img, pt1, pt2, color, thickness)

def draw_elements(img, elements):
    """Draw all saved element rectangles"""
    colors = {
        'Departure': (255, 0, 0),      # Blue
        'Arrival': (0, 255, 0),        # Green
        'TrainNo': (0, 0, 255),        # Red
        'Class': (255, 255, 0),        # Cyan
        'Time': (255, 0, 255),         # Purple
        'Price': (0, 255, 255),        # Yellow
        'Other': (128, 128, 128)       # Gray
    }

    for elem in elements:
        color = colors.get(elem['type'], (128, 128, 128))
        draw_rectangle(img, (elem['x1'], elem['y1']), (elem['x2'], elem['y2']), color, 2)
        # Display element type
        cv2.putText(img, elem['type'], (elem['x1'], elem['y1'] - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

def mouse_callback(event, x, y, flags, param):
    """Mouse event handler"""
    global ref_point, dragging, current_element, img_copy, img_original

    if event == cv2.EVENT_LBUTTONDOWN:
        # Start dragging
        ref_point = [(x, y)]
        dragging = True

    elif event == cv2.EVENT_MOUSEMOVE:
        # Dragging, draw rectangle in real-time
        if dragging and ref_point:
            img_copy = img_original.copy()
            draw_elements(img_copy, elements)
            draw_rectangle(img_copy, ref_point[0], (x, y))
            # Display current selected element type
            cv2.putText(img_copy, f'Current type: {element_types[current_type_index]}',
                        (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)

    elif event == cv2.EVENT_LBUTTONUP:
        # End dragging, save rectangle
        if dragging and ref_point:
            ref_point.append((x, y))
            dragging = False

            # Ensure correct coordinate order (top-left to bottom-right)
            x1, y1 = ref_point[0]
            x2, y2 = ref_point[1]
            if x1 > x2:
                x1, x2 = x2, x1
            if y1 > y2:
                y1, y2 = y2, y1

            # Save element
            element = {
                'type': element_types[current_type_index],
                'x1': x1,
                'y1': y1,
                'x2': x2,
                'y2': y2,
                'width': x2 - x1,
                'height': y2 - y1
            }
            elements.append(element)
            print(f"Added element: {element['type']} - Coordinates: ({x1}, {y1}) to ({x2}, {y2})")

            # Update display
            img_copy = img_original.copy()
            draw_elements(img_copy, elements)
            cv2.putText(img_copy, f'Current type: {element_types[current_type_index]}',
                        (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)

def pixel_to_ratio_coordinates(pixel_coords, image_width, image_height):
    """将像素坐标转换为比例坐标"""
    ratio_coords = []
    for elem in pixel_coords:
        ratio_elem = {
            'type': elem['type'],
            'x1': elem['x1'] / image_width,
            'y1': elem['y1'] / image_height,
            'x2': elem['x2'] / image_width,
            'y2': elem['y2'] / image_height,
            'width_ratio': elem['width'] / image_width,
            'height_ratio': elem['height'] / image_height
        }
        ratio_coords.append(ratio_elem)
    return ratio_coords

def interactive_element_selection(image_path, output_path='output'):
    """Interactive element selection"""
    global img_original, img_copy, elements, current_type_index

    # Create output directory
    Path(output_path).mkdir(exist_ok=True)

    # Reset global variables
    elements = []
    current_type_index = 0

    # Read image
    img_original = cv2.imread(image_path)
    if img_original is None:
        print(f"Cannot read image: {image_path}")
        return

    img_copy = img_original.copy()
    height, width = img_original.shape[:2]

    # Create window
    cv2.namedWindow("Train Ticket Element Selection", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Train Ticket Element Selection", 800, 600)
    cv2.setMouseCallback("Train Ticket Element Selection", mouse_callback)

    print("\nInteractive Train Ticket Element Selection Tool")
    print("=" * 50)
    print("Instructions:")
    print("1. Click and drag to select element region")
    print("2. Press number keys 1-7 to select element type:")
    for i, elem_type in enumerate(element_types):
        print(f"   {i+1}. {elem_type}")
    print("3. Press 'd' to delete last element")
    print("4. Press 'c' to clear all elements")
    print("5. Press 's' to save results (both pixel and ratio coordinates)")
    print("6. Press 'q' to quit")
    print("=" * 50)

    # Display initial information
    cv2.putText(img_copy, f'Current type: {element_types[current_type_index]}',
                (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)

    while True:
        cv2.imshow("Train Ticket Element Selection", img_copy)
        key = cv2.waitKey(1) & 0xFF

        # Select element type
        if key in [ord(str(i+1)) for i in range(len(element_types))]:
            current_type_index = int(chr(key)) - 1
            print(f"Current element type: {element_types[current_type_index]}")
            # Update display
            img_copy = img_original.copy()
            draw_elements(img_copy, elements)
            cv2.putText(img_copy, f'Current type: {element_types[current_type_index]}',
                        (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)

        # Delete last element
        elif key == ord('d'):
            if elements:
                removed = elements.pop()
                print(f"Deleted element: {removed['type']}")
                # Update display
                img_copy = img_original.copy()
                draw_elements(img_copy, elements)
                cv2.putText(img_copy, f'Current type: {element_types[current_type_index]}',
                            (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)

        # Clear all elements
        elif key == ord('c'):
            elements.clear()
            print("Cleared all elements")
            # Update display
            img_copy = img_original.copy()
            cv2.putText(img_copy, f'Current type: {element_types[current_type_index]}',
                        (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)

        # Save results
        elif key == ord('s'):
            if not elements:
                print("No elements to save")
                continue

            # Save pixel coordinates to JSON file
            pixel_json_file = f'{output_path}/element_coordinates.json'
            with open(pixel_json_file, 'w', encoding='utf-8') as f:
                json.dump(elements, f, ensure_ascii=False, indent=2)

            # Convert to ratio coordinates
            ratio_coords = pixel_to_ratio_coordinates(elements, width, height)

            # Save ratio coordinates to JSON file
            ratio_json_file = f'{output_path}/element_ratio_coordinates.json'
            with open(ratio_json_file, 'w', encoding='utf-8') as f:
                json.dump(ratio_coords, f, ensure_ascii=False, indent=2)

            # Save marked image
            marked_img = img_original.copy()
            draw_elements(marked_img, elements)
            cv2.imwrite(f'{output_path}/marked_elements.jpg', marked_img)

            print(f"Saved {len(elements)} pixel elements to {pixel_json_file}")
            print(f"Saved {len(ratio_coords)} ratio elements to {ratio_json_file}")
            print(f"Saved marked image to {output_path}/marked_elements.jpg")

        # Quit
        elif key == ord('q'):
            break

    cv2.destroyAllWindows()
    return elements

def main():
    """Main function"""
    # Default to use rotated ticket image
    current_dir = Path(__file__).parent
    image_path = str(current_dir / 'output' /'refer_ticket.jpg')

    # Check if image exists
    if not Path(image_path).exists():
        # If not, find image files in current directory

        print("No image files found, please place images in the same directory or output directory")
        return


    print(f"Using image: {image_path}")
    interactive_element_selection(image_path)

if __name__ == "__main__":
    main()