conveyor_visualizer.py

conveyor_visualizer.py

import cv2
import numpy as np

#create window
#define pulles


# Configuration Variables
WINDOW_SIZE = (2000, 4096)
width = WINDOW_SIZE[1]
height = WINDOW_SIZE[0]
GRID_SIZE = 50
INTERPOLATED_POINTS = 600     # Number of interpolated points
CIRCLE_RADIUS = 100              # Radius of the green circles
LINE_COLOR = (255, 0, 0)        # Color of the line (Blue in BGR format)
CIRCLE_COLOR = (0, 255, 0)      # Color of the circles (Green in BGR format)
LINE_THICKNESS = 8              # Thickness of the line
INTERPOLATION_KEY = 'i'         # Key to trigger interpolation
REMOVE_POINTS_KEY = 't'         # Key to trigger removing middle points
EXIT_KEY = 'q'                  # Key to exit the program
TOGGLE_MODE_KEY = 13            # Enter key to switch between modes (ASCII 13)
SMOOTHING_WINDOW_LENGTH = 35
DRIVE_STATUS = False

fourcc = cv2.VideoWriter_fourcc(*'XVID')  # Codec for .avi files
out = cv2.VideoWriter('/home/smeschke/Desktop/videos/a1.avi', fourcc, 60.0, (width, height))

# Global variables
draw_circle_mode = True  # Initially in circle drawing mode
points = []  # To store points for the blue line
green_circles = []  # To store green circle positions and radii, and drive status
takeup_lines = []


def handle_mouse_events(event, x, y, flags, param):
    global draw_circle_mode, points, green_circles, takeup_lines

    # Define a callback function for mouse events
    if draw_circle_mode and event == cv2.EVENT_LBUTTONDOWN:
        if CIRCLE_RADIUS == 5:
            takeup_lines.append((x,y))
        else:    
            # Draw a green circle at the clicked position
            cv2.circle(img, (x, y), CIRCLE_RADIUS, CIRCLE_COLOR, -1)
            green_circles.append((x, y, CIRCLE_RADIUS, DRIVE_STATUS))  # Non-driven by default
            
    elif draw_circle_mode and event == cv2.EVENT_RBUTTONDOWN:
        # Toggle drive status of the circle if right-clicked
        for i, circle in enumerate(green_circles):
            center, radius, is_drive = circle[:2], circle[2], circle[3]
            if np.linalg.norm(np.array(center) - np.array((x, y))) <= radius:
                green_circles[i] = (center[0], center[1], radius, not is_drive)  # Toggle is_drive
                break
        
    elif not draw_circle_mode and event == cv2.EVENT_LBUTTONDOWN:
        if len(points) == 0:
            points.append((x, y))
        else:
            # Draw a continuous blue line from the last point to the current point
            cv2.line(img, points[-1], (x, y), LINE_COLOR, LINE_THICKNESS)
            points.append((x, y))  # Store the current point


def draw_grid(img, width, height, grid_size):
    """
    Draws a grid over the image.
    
    :param img: The image to draw on.
    :param window_size: Tuple (width, height) for the image size.
    :param grid_size: The size of each grid cell.
    """
    for x in range(0, width, grid_size):
        cv2.line(img, (x, 0), (x, height), (200, 200, 200), 1
                 )  # Vertical grid lines
    for y in range(0, height, grid_size):
        cv2.line(img, (0, y), (width, y), (200, 200, 200), 1)  # Horizontal grid lines



# Function to interpolate and create exactly INTERPOLATED_POINTS equally spaced points along the line, including closing the loop
def generate_interpolated_points(points, num_points=INTERPOLATED_POINTS):
    interpolated_points = []
    total_distance = 0
    
    # Calculate the total distance of the polyline (sum of segments), including the loop closure
    for i in range(len(points)):
        p1 = np.array(points[i])
        p2 = np.array(points[(i + 1) % len(points)])  # Connect the last point to the first point
        total_distance += np.linalg.norm(p2 - p1)

    # Create equally spaced points along the polyline
    segment_distance = total_distance / num_points
    current_distance = 0
    
    # Generate points
    for i in range(len(points)):
        p1 = np.array(points[i])
        p2 = np.array(points[(i + 1) % len(points)])  # Connect the last point to the first point
        segment_length = np.linalg.norm(p2 - p1)
        
        while current_distance <= segment_length and len(interpolated_points) < num_points:
            t = current_distance / segment_length
            new_point = (1 - t) * p1 + t * p2
            interpolated_points.append(tuple(map(int, new_point)))  # Convert to int tuple
            current_distance += segment_distance
        
        current_distance -= segment_length

    return interpolated_points


# Function to redraw the points as dots instead of lines
def redraw(interpolatedPoints):
    # Clear the image and redraw circles
    img[:] = 0  # Reset the image to white
    for circle in green_circles:
        cv2.circle(img, circle[:2], circle[2], CIRCLE_COLOR, -1)  # Redraw green circles
    
    # Draw points instead of lines
    for point in interpolatedPoints:
        cv2.circle(img, point, 2, LINE_COLOR, -1)  # Draw small circle (point) at each point's location

    # Close the loop by drawing the last point to the first as dots
    if len(interpolatedPoints) > 2:
        cv2.circle(img, points[-1], 2, LINE_COLOR, -1)  # Draw dot for the last point
        cv2.circle(img, points[0], 2, LINE_COLOR, -1)   # Draw dot for the first point to complete the loop


# Function to interpolate points, update them, and redraw the line
def interpolate_and_redraw():
    global points
    points = generate_interpolated_points(points, INTERPOLATED_POINTS)  # Interpolate to 10,000 points
    redraw_points()


# Function to interpolate points, update them, and redraw the line
def interpolate():
    global points
    points = generate_interpolated_points(points, INTERPOLATED_POINTS)  # Interpolate to 10,000 points
    return(points)


# Create a blank white image
img = 0 * np.ones(WINDOW_SIZE + (3,), dtype=np.uint8)
cv2.putText(img, "Mode: Circle Drawing", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

# Create a named window
cv2.namedWindow('Image')

# Set the mouse callback function
cv2.setMouseCallback('Image', handle_mouse_events)

        


while True:
    
    # Display the image
    cv2.imshow('Image', img)
    #out.write(img)
    draw_grid(img, width, height, GRID_SIZE)
    
    # Check for key press
    key = cv2.waitKey(1) & 0xFF
    
    # Press EXIT_KEY to exit
    if key == ord(EXIT_KEY):
        break
    if key == 115:  # Up arrow key
        CIRCLE_RADIUS += 5
        print(f"Circle radius increased to: {CIRCLE_RADIUS}")
    if key == 116:  # Down arrow key
        CIRCLE_RADIUS = max(5, CIRCLE_RADIUS - 5)  # Ensure radius doesn't go below 5
        print(f"Circle radius decreased to: {CIRCLE_RADIUS}")

    if key == 117:  # Up arrow key
        DRIVE_STATUS = False
        print("Drive status: false")
    if key == 118:  # Down arrow key
        DRIVE_STATUS = True
        print("Drive status: true")
    
    # Press TOGGLE_MODE_KEY (Enter) to switch between circle and line drawing mode
    elif key == TOGGLE_MODE_KEY:  # Enter key is ASCII 13
        if draw_circle_mode:
            draw_circle_mode = False
            cv2.putText(img, "Mode: Line Drawing", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

            

            print("Switched to drawing continuous blue line.")
        else:
            # Close the loop by drawing the last point to the first
            if len(points) > 1:
                cv2.line(img, points[-1], points[0], LINE_COLOR, LINE_THICKNESS)
            draw_circle_mode = True
            print("Switched to drawing green circles.")
    
    # Press INTERPOLATION_KEY to interpolate the points and redraw the smooth line
    elif key == ord(INTERPOLATION_KEY):
        points = interpolate()
        redraw(points)



# Close the window
cv2.destroyAllWindows()

from scipy.signal import savgol_filter
import numpy as np

def smooth_points(points, window_length=9, polyorder=2):
    """
    Smooths a list of 2D points using Savitzky-Golay filter.

    :param points: List of tuples [(x1, y1), (x2, y2), ...]
    :param window_length: The length of the filter window (must be odd).
    :param polyorder: The order of the polynomial used to fit the samples.
    :return: List of smoothed points [(x1_smooth, y1_smooth), (x2_smooth, y2_smooth), ...]
    """
    # Separate the x and y coordinates
    x_coords = np.array([p[0] for p in points])
    y_coords = np.array([p[1] for p in points])

    # Apply the Savitzky-Golay filter to smooth x and y coordinates independently
    x_smooth = savgol_filter(x_coords, window_length, polyorder)
    y_smooth = savgol_filter(y_coords, window_length, polyorder)

    # Combine the smoothed x and y coordinates back into a list of points
    smoothed_points = list(zip(x_smooth, y_smooth))

    return smoothed_points

def get_10_points(points, start_index):
    """
    Get 10 consecutive points from the list, wrapping around if necessary.
    
    :param points: List of points [(x1, y1), (x2, y2), ...]
    :param start_index: The index from which to start retrieving points.
    :return: List of 10 consecutive points.
    """
    total_points = len(points)
    selected_points = []

    for i in range(10):
        index = (start_index + i) % total_points  # Wrap around if end of list is reached
        selected_points.append(points[index])
    
    return selected_points


def is_point_inside_any_circle(point, circles):
    """
    Checks if a point is inside or touching any circle from a list of circles.

    :param point: Tuple (x, y) representing the point's coordinates.
    :param circles: List of tuples (x, y, r) where (x, y) are the circle center coordinates and r is the radius.
    :return: True if the point is inside or touching any of the circles, False otherwise.
    """
    for circle in circles:
        circle_center = np.array(circle[:2])  # Extract the circle's center (x, y)
        radius = circle[2]  # Extract the circle's radius

        # Calculate the distance between the point and the circle center
        distance = np.linalg.norm(np.array(point) - circle_center)

        # Check if the distance is less than or equal to the radius
        if distance <= radius:
            return True  # Point is inside or touching this circle
    return False  # Point is not inside any circle

index = -SMOOTHING_WINDOW_LENGTH


import math

def draw_rotating_line(img, circle, angle):
    """
    Draws a rotating line on a pulley to show that it is driven.
    
    :param img: The image where the line will be drawn.
    :param circle: Tuple (x, y, r, is_drive), where (x, y) is the circle center, r is the radius.
    :param angle: The current angle in radians that the line should rotate to.
    """
    center = np.array(circle[:2])  # Circle center (x, y)
    radius = circle[2]  # Circle radius

    # Calculate the endpoint of the line based on the angle
    end_x = int(center[0] + radius * math.cos(angle))
    end_y = int(center[1] + radius * math.sin(angle))
    endpoint = (end_x, end_y)

    # Draw the rotating line from the center to the calculated endpoint
    cv2.line(img, tuple(center), endpoint, (0, 0, 255), 6)  # Red line for driven pulleys

def hsv_to_bgr(h, s, v):
    """
    Converts HSV color to BGR for OpenCV display.
    :param h: Hue value (0-180 for OpenCV).
    :param s: Saturation value (0-255).
    :param v: Value (brightness) (0-255).
    :return: BGR tuple.
    """
    hsv_color = np.uint8([[[h, s, v]]])  # Create an HSV image
    bgr_color = cv2.cvtColor(hsv_color, cv2.COLOR_HSV2BGR)  # Convert to BGR
    return tuple(int(c) for c in bgr_color[0][0])  # Return as tuple


angle = 0
while True:
    
    img[:] = 0  # Reset the image to white
    angle += .1

    
    #cv2.line(img, takeup_lines[0], takeup_lines[1], (123,123,123),5)

    

    # Draw pulleys and spinning line for driven pulleys
    for circle in green_circles:
        center, radius, is_drive = circle[:2], circle[2], circle[3]

        # Draw the circle (use red for driven pulleys)
        
        

        # If it's a driven pulley, draw the rotating line
        if is_drive:
            cv2.circle(img, center, radius, (255,255,0), -1)  # Draw the pulley
            draw_rotating_line(img, circle, angle)
        else:
            cv2.circle(img, center, radius, (0,255,255), -1)  # Draw the pulley

    currentPoints = []
    currentIndex = index
    for ppp in range(SMOOTHING_WINDOW_LENGTH):
        currentPoints.append(points[index+ppp])
        

    smoothedPoints = smooth_points(currentPoints)
        

    for point in points:
        point = tuple(np.array(point, int))
        cv2.circle(img, point, 8, (255,255,255), -1)  # Draw small circle (point) at each point's location
        


    # Draw points with rainbow colors
    for i, point in enumerate(smoothedPoints):
        point = tuple(np.array(point, int))
        
        # Calculate the hue value (0-180 for OpenCV)
        hue = int((i / len(smoothedPoints)) * 180)  # Scale index to hue (0-180 range for OpenCV)
        
        # Convert HSV to BGR color (full saturation and brightness)
        color = hsv_to_bgr(hue, 255, 255)  # Full saturation and brightness for vibrant colors
        
        # Draw the point with the rainbow color
        cv2.circle(img, point, 10, color, -1)  # Draw small circle at each point's location


    # reassign the points if it is not insdie the circle
    for ppp in range(SMOOTHING_WINDOW_LENGTH):
        smoothedPoint = smoothedPoints[ppp]
        pointInside = is_point_inside_any_circle(smoothedPoint, green_circles)
        if  not pointInside:
            points[index+ppp] = smoothedPoints[ppp]

    

    # Increment the index and wrap around if needed
    index += 1
    if index > len(points)-SMOOTHING_WINDOW_LENGTH: index = -SMOOTHING_WINDOW_LENGTH
    
        
    cv2.imshow('img', img)
    out.write(img)
    k = cv2.waitKey(1)

    # Press EXIT_KEY to exit
    if key == 27:
        break
    



#get pullies
#show tangent lines
#pick lines
#fill in circle parts
#draw scrapers
#draw idlers
#draw returns
#draw pans
#draw skirtboard
#draw sensors