detect.py

from json import load
import cv2
import numpy as np

# MJPEG stream from Raspberry Pi
stream_url = "http://192.168.1.200:8080"  # Replace with actual IP of Pi
cap = cv2.VideoCapture(stream_url)

# Pre-defined regions for turnout status (example coordinates for the "closed" and "open" regions)
turnout_open_template = cv2.imread("turnouts/open_template.jpg", cv2.IMREAD_GRAYSCALE)
turnout_closed_template = cv2.imread(
    "turnouts/closed_template.jpg", cv2.IMREAD_GRAYSCALE
)


def extract_features(image):
    """
    Extract color histogram features for train recognition
    """

    if image is None:
        raise ValueError("Image is None. Check the file path.")
    
    image = cv2.resize(image, (100, 100))  # Resize for consistent comparison
    hist = cv2.calcHist([image], [0, 1, 2], None, [8, 8, 8], [0, 256] * 3)
    return cv2.normalize(hist, hist).flatten()


def identify_train(frame, known_trains):
    """
    Compare each frame with known train templates to identify the train
    """
    f = extract_features(frame)
    scores = {
        name: cv2.compareHist(f, vec, cv2.HISTCMP_CORREL)
        for name, vec in known_trains.items()
    }
    # return max(scores, key=scores.get) if max(scores.values()) > 0.8 else "unknown"


def detect_tracks(frame) -> list:
    """
    Detect model railway tracks (straight and curved) by identifying edge pairs
    that resemble rail tracks.
    Returns a list of tuples: [(pt1, pt2), ...] for lines
    """
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    sleeper_template = cv2.imread("turnouts/sleeper.jpg", cv2.IMREAD_GRAYSCALE)

    # Match template
    res = cv2.matchTemplate(frame_gray, sleeper_template, cv2.TM_CCOEFF_NORMED)
    threshold = 0.6
    loc = np.where(res >= threshold)

    # Group close detections (sleepers are evenly spaced)
    detections = list(zip(*loc[::-1]))
    grouped = []
    for pt in detections:
        added = False
        for group in grouped:
            if abs(pt[0] - group[0][0]) < 10 and abs(pt[1] - group[0][1]) < 10:
                group.append(pt)
                added = True
                break
        if not added:
            grouped.append([pt])

    # Filter groups with enough detections (e.g. >3 sleepers in a line)
    tracks = []
    for group in grouped:
        if len(group) >= 3:
            # Draw line between first and last
            sorted_pts = sorted(group, key=lambda x: x[1])  # sort vertically
            start = sorted_pts[0]
            end = sorted_pts[-1]
            tracks.append((start, end))

    return tracks


def turnout_status(frame, region):
    """
    Identify whether the turnout is open or closed based on a region of interest (ROI)
    """
    x, y, w, h = region
    roi = frame[y : y + h, x : x + w]

    # Match the turnout template with the region of interest (ROI)
    match_open = cv2.matchTemplate(roi, turnout_open_template, cv2.TM_CCOEFF_NORMED)
    match_closed = cv2.matchTemplate(roi, turnout_closed_template, cv2.TM_CCOEFF_NORMED)

    open_score = match_open.max()
    closed_score = match_closed.max()

    if open_score > closed_score:
        return "open"
    else:
        return "closed"
    
def load_known_trains():
    try:
        return {
            "express": extract_features(cv2.imread("trains/express.jpg")),
            "freight": extract_features(cv2.imread("trains/freight.jpg")),
        }
    except ValueError as e:
        print("Failed to load known train images:", e)
        return {}


def process_frame(frame):
    """
    Process the video frame for train detection, track detection, and turnout status
    """
    # Example of known train templates you need to train these images manually)
    # known_trains = load_known_trains()

    # # Detect trains
    # train = identify_train(frame, known_trains)

    # Detect tracks (lines)
    lines = detect_tracks(frame)

    # Detect turnout status (assuming the region of interest is known)
    # turnout_region = (
    #     100,
    #     200,
    #     50,
    #     50,
    # )  # Example (x, y, width, height) of turnout region
    # turnout = turnout_status(frame, turnout_region)

    return lines


while True:
    ret, frame = cap.read()
    if not ret:
        print("Failed to capture frame")
        break

    # Process the frame
    tracks = detect_tracks(frame)

    # Display results
    # print(f"Train Detected: {train}")
    # print(f"Turnout Status: {turnout}")
    print(f"Track Lines Detected: {len(tracks)}")

    # 1. Draw train detection (Example bounding box for detected trains)
    # if train != "unknown":
    #     cv2.putText(
    #         frame,
    #         f"Train: {train}",
    #         (30, 30),
    #         cv2.FONT_HERSHEY_SIMPLEX,
    #         1,
    #         (255, 0, 0),
    #         2,
    #     )

    # 2. Draw track lines
    if tracks is not None:
       for (x1, y1), (x2, y2) in tracks:
            cv2.line(frame, (x1, y1), (x2, y2), (255, 0, 255), 2)  # Magenta centerline

    # 3. Draw turnout status (assuming a defined ROI is found)
    # x, y, w, h = (
    #     100,
    #     200,
    #     50,
    #     50,
    # )
    # color = (
    #     (0, 0, 255) if turnout == "closed" else (0, 255, 0)
    # )  # Red for closed, Green for open
    # cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
    # cv2.putText(
    #     frame,
    #     f"Turnout: {turnout}",
    #     (x, y - 10),
    #     cv2.FONT_HERSHEY_SIMPLEX,
    #     0.6,
    #     color,
    #     2,
    # )

    # Show the frame with the detections
    cv2.namedWindow("Train Detection", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Train Detection", 800, 600)
    cv2.imshow("Train Detection", frame)

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

cap.release()
cv2.destroyAllWindows()