Untitled

 avatar
unknown
plain_text
5 months ago
4.6 kB
3
Indexable
import cv2
import numpy as np
import time
from picamera2 import Picamera2

# Initialize the camera
picam2 = Picamera2()
picam2.configure(picam2.create_preview_configuration(main={"format": 'XRGB8888', "size": (640, 480)}))
picam2.start()

def detect_pingpong(frame, roi, detection_started_time, detection_delay):
    """
    Detect ping pong balls within a specific region of interest (ROI) and classify them based on the mask.
    """
    x1, y1, x2, y2 = roi
    cropped_frame = frame[y1:y2, x1:x2]

    # Convert to HSV color space for color filtering
    hsv = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2HSV)

    # Define the HSV range for the ping pong ball color (orange)
    lower_orange = np.array([10, 100, 100])  # Adjust as needed
    upper_orange = np.array([25, 255, 255])

    # Create a mask for the orange color
    mask = cv2.inRange(hsv, lower_orange, upper_orange)

    # Apply morphological operations to clean up the mask
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    # Find contours in the mask
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    for contour in contours:
        # Approximate the contour to a circle
        ((x, y), radius) = cv2.minEnclosingCircle(contour)

        # Filter based on size (radius range for ping pong balls)
        if 20 < radius < 50:  # Adjust size range as needed
            # Check if detection timer has started
            if detection_started_time is None:
                detection_started_time = time.time()  # Start the timer
            elif time.time() - detection_started_time >= detection_delay:
                # Extract the ball region from the mask
                ball_mask = mask[int(y - radius):int(y + radius), int(x - radius):int(x + radius)]

                # Check if the ball_mask size is valid
                if ball_mask.size == 0:
                    print("Ball mask size is zero; skipping this contour.")
                    continue  # Skip if the mask size is zero

                # Debugging: Print pixel counts
                total_pixels = ball_mask.size
                white_pixels = cv2.countNonZero(ball_mask)
                black_pixels = total_pixels - white_pixels
                print(f"Total Pixels: {total_pixels}, White Pixels: {white_pixels}, Black Pixels: {black_pixels}")

                # Check for black spots in the mask
                black_pixel_percentage = (black_pixels / total_pixels) * 100

                if black_pixel_percentage > 5:  # Threshold for detecting defects
                    label = "Defect"
                    color = (0, 0, 255)  # Red for defective
                else:
                    label = "Perfect (0% Defect)"
                    color = (0, 255, 0)  # Green for non-defective

                # Draw a bounding box tightly around the ball
                top_left = (int(x - radius) + x1, int(y - radius) + y1)  # Map back to full frame
                bottom_right = (int(x + radius) + x1, int(y + radius) + y1)  # Map back to full frame
                cv2.rectangle(frame, top_left, bottom_right, color, 2)

                # Annotate the frame with the defect label
                cv2.putText(frame, label, (top_left[0], top_left[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

        else:
            # Reset the timer if the ball leaves the ROI
            detection_started_time = None

    # Draw the ROI boundary on the original frame
    cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 255, 0), 2)

    return frame, mask, detection_started_time


print("Press 'q' to exit...")

# Define the ROI based on your coordinates
roi = (202, 16, 476, 361)  # (x1, y1, x2, y2)
detection_delay = 1.0  # Delay in seconds
detection_started_time = None  # Timer for detection delay

while True:
    # Capture frame from Picamera2
    frame = picam2.capture_array()

    # Detect and classify ping pong balls within the ROI
    processed_frame, mask, detection_started_time = detect_pingpong(
        frame, roi, detection_started_time, detection_delay
    )

    # Show the processed frame and color mask
    cv2.imshow("Defect Detection", processed_frame)
    cv2.imshow("Color Mask", mask)

    # Press 'q' to quit
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# Release resources
cv2.destroyAllWindows()
picam2.stop()
Editor is loading...
Leave a Comment