Add classical filter implementation

classical-filter/classical-filter.cpp

@@ -0,0 +1,132 @@
+#include <opencv2/opencv.hpp>
+#include <iostream>
+#include <cmath>
+
+using namespace cv;
+using namespace std;
+
+// Measures how similar two numbers are
+double sim(double a, double b) {
+    return min(a, b) / (a + b);
+}
+
+// Calculates the squared Euclidean distance between two points
+double distSq(Point2f a, Point2f b) {
+    return pow(a.x - b.x, 2) + pow(a.y - b.y, 2);
+}
+
+// Calculates the angle between two points in degrees
+double angle(Point2f a, Point2f b) {
+    if (a.x == b.x) return 90;
+    Point2f right = (a.x > b.x) ? a : b;
+    Point2f left = (a.x > b.x) ? b : a;
+    return atan((right.y - left.y) / (right.x - left.x)) * 180 / CV_PI;
+}
+
+// Applies a mask to detect color-specific contours in the frame
+vector<vector<Point>> getContours(Mat& frame, const string& color) {
+    Mat frameHSV, mask1, mask2, frameThreshold;
+
+    cvtColor(frame, frameHSV, COLOR_BGR2HSV);
+
+    if (color == "red") {
+        inRange(frameHSV, Scalar(0, 70, 50), Scalar(20, 255, 255), mask1);
+        inRange(frameHSV, Scalar(170, 70, 50), Scalar(230, 255, 255), mask2);
+    } else {
+        inRange(frameHSV, Scalar(90, 70, 50), Scalar(120, 255, 255), mask1);
+        inRange(frameHSV, Scalar(170, 70, 50), Scalar(200, 255, 255), mask2);
+    }
+    frameThreshold = mask1 | mask2;
+
+    erode(frameThreshold, frameThreshold, getStructuringElement(MORPH_ELLIPSE, Size(5, 5)));
+    dilate(frameThreshold, frameThreshold, getStructuringElement(MORPH_ELLIPSE, Size(5, 5)));
+    dilate(frameThreshold, frameThreshold, getStructuringElement(MORPH_ELLIPSE, Size(5, 5)));
+    erode(frameThreshold, frameThreshold, getStructuringElement(MORPH_ELLIPSE, Size(5, 5)));
+
+    vector<vector<Point>> contours;
+    findContours(frameThreshold, contours, RETR_EXTERNAL, CHAIN_APPROX_NONE);
+
+    return contours;
+}
+
+// Draws the bounding box centers and processes relationships between them
+Mat drawCenters(Mat frame, const string& color) {
+    auto contours = getContours(frame, color);
+    vector<RotatedRect> bboxes;
+
+    for (auto& contour : contours) {
+        RotatedRect bbox = minAreaRect(contour);
+        bboxes.push_back(bbox);
+
+        Point2f bboxPoints[4];
+        bbox.points(bboxPoints);
+        for (int j = 0; j < 4; j++) {
+            line(frame, bboxPoints[j], bboxPoints[(j + 1) % 4], Scalar(0, 255, 0), 2);
+        }
+    }
+
+    int thresh = 20;
+    double widthSimThresh = 0.1, lengthSimThresh = 0.3, yThresh = 15, angleThresh = 15;
+
+    for (size_t i = 0; i < bboxes.size(); i++) {
+        auto& bbox1 = bboxes[i];
+        Point2f center1 = bbox1.center;
+        double width1 = bbox1.size.width;
+        double length1 = bbox1.size.height;
+        double angle1 = bbox1.angle;
+
+        if (max(length1, width1) < thresh) continue;
+
+        for (size_t j = i + 1; j < bboxes.size(); j++) {
+            auto& bbox2 = bboxes[j];
+            Point2f center2 = bbox2.center;
+            double width2 = bbox2.size.width;
+            double length2 = bbox2.size.height;
+            double angle2 = bbox2.angle;
+
+            if (max(length2, width2) < thresh) continue;
+
+            double angleDiff = abs(angle1 - angle2);
+            double yDiff = abs(center1.y - center2.y);
+
+            if (sim(width1, width2) > widthSimThresh &&
+                sim(length1, length2) > lengthSimThresh &&
+                yDiff < yThresh &&
+                (angleDiff < angleThresh || angleDiff > 180 - angleThresh)) {
+                Point centerMid((center1.x + center2.x) / 2, (center1.y + center2.y) / 2);
+                circle(frame, centerMid, 10, Scalar(255, 0, 255), -1);
+            }
+        }
+    }
+
+    return frame;
+}
+
+int main() {
+    VideoCapture cap("tests/test1.mp4");
+
+    if (!cap.isOpened()) {
+        cout << "Error opening video file" << endl;
+        return -1;
+    }
+
+    while (cap.isOpened()) {
+        Mat frame;
+        cap >> frame;
+
+        if (frame.empty()) {
+            cout << "End of video stream." << endl;
+            break;
+        }
+
+        frame = drawCenters(frame, "blue");
+
+        imshow("Frame", frame);
+
+        if (waitKey(1) == 'q') break;
+    }
+
+    cap.release();
+    destroyAllWindows();
+    return 0;
+}

classical-filter/classical-filter.py

@@ -0,0 +1,183 @@
+# References:
+# https://docs.opencv.org/4.x/dd/d43/tutorial_py_video_display.html
+# https://docs.opencv.org/4.x/d9/dc8/tutorial_py_trackbar.html
+# https://docs.opencv.org/4.x/d7/d4d/tutorial_py_thresholding.html
+# https://www.opencv-srf.com/2010/09/object-detection-using-color-seperation.html
+# https://docs.opencv.org/4.5.4/da/d97/tutorial_threshold_inRange.html
+# https://docs.opencv.org/4.5.4/db/df6/tutorial_erosion_dilatation.html
+# https://docs.opencv.org/3.4/d4/d73/tutorial_py_contours_begin.html
+# https://docs.opencv.org/4.x/dc/da5/tutorial_py_drawing_functions.html
+# https://stackoverflow.com/questions/38064777/use-waitkey-in-order-pause-and-play-video
+# https://docs.opencv.org/4.x/de/d62/tutorial_bounding_rotated_ellipses.html
+# https://www.geeksforgeeks.org/python-opencv-cv2-puttext-method/
+
+import cv2 as cv
+import numpy as np
+import math
+import time
+
+color = 'blue'
+
+# Measures how similar two numbers are
+def sim(a, b):
+    return min(a, b) / (a + b)
+
+def dist_sq(a, b):
+    return (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2
+
+def angle(a, b):
+    if a[0] == b[0]:
+        return 90
+    right = a
+    left = b
+    if b[0] > a[0]:
+        right = b
+        left = a
+    return math.degrees(math.atan((right[1] - left[1]) / (right[0] - left[0])))
+
+def get_contours(frame, color):
+    # Apply a red mask to image, apply morphological opening/closing, and find contours of contiguous red areas
+    frame_HSV = None
+
+    mask1, mask2, frame_threshold = None, None, None
+    if color == 'red':
+        frame_HSV = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
+        mask1 = cv.inRange(frame_HSV, (0, 70, 50), (20, 255, 255))
+        mask2 = cv.inRange(frame_HSV, (170, 70, 50), (230, 255, 255))
+        frame_threshold = mask1 | mask2
+    else:
+        frame_HSV = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
+        mask1 = cv.inRange(frame_HSV, (90, 100, 100), (115, 255, 255))
+        mask2 = cv.inRange(frame_HSV, (115, 100, 100), (135, 255, 255))
+        frame_threshold = mask1 | mask2
+
+    frame_threshold = cv.erode(frame_threshold, cv.getStructuringElement(cv.MORPH_ELLIPSE, (5, 5)))
+    frame_threshold = cv.dilate(frame_threshold, cv.getStructuringElement(cv.MORPH_ELLIPSE, (5, 5)))
+
+    frame_threshold = cv.dilate(frame_threshold, cv.getStructuringElement(cv.MORPH_ELLIPSE, (5, 5)))
+    frame_threshold = cv.erode(frame_threshold, cv.getStructuringElement(cv.MORPH_ELLIPSE, (5, 5)))
+
+    contours, _ = cv.findContours(frame_threshold, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)
+
+    return contours
+
+def draw_centers(frame, color):
+    contours = get_contours(frame, color)
+
+    # Compute rotated bounding box for each contour and store in `bboxes`
+    bboxes = []
+    for contour in contours:
+        bbox = cv.minAreaRect(contour)
+        bboxes.append(bbox)
+        bbox_points = cv.boxPoints(bbox)
+        bbox_points = np.intp(bbox_points)
+        frame = cv.drawContours(frame, [bbox_points], -1, (0, 255, 0), 2)
+
+    thresh = 20
+    width_sim_thresh, length_sim_thresh, y_thresh, angle_thresh = 0.1, 0.3, 0.15, 15
+    for i in range(len(bboxes)):
+        bbox1 = bboxes[i]
+        center1 = (bbox1[0][0], bbox1[0][1])
+        width1, length1, angle1 = bbox1[1][0], bbox1[1][1], bbox1[2]
+
+        vert1 = cv.boxPoints(bbox1)
+
+        longest = (0, 1)
+        bl = vert1[0]
+        tl = vert1[1]
+        br = vert1[3]
+        width1 = math.sqrt(dist_sq(vert1[0], vert1[3]))
+        if dist_sq(bl, br) > dist_sq(bl, tl):
+            longest = (0, 3)
+            width1 = math.sqrt(dist_sq(vert1[0], vert1[1]))
+        length1 = math.sqrt(dist_sq(vert1[longest[0]], vert1[longest[1]]))
+        angle1 = angle(vert1[longest[0]], vert1[longest[1]])
+
+        # If the bounding box is too small, skip
+        if max(length1, width1) < thresh:
+            continue
+
+        # how the heck is `minAreaRect` defining the angle
+        # if angle1 > 45:
+        #     angle1 = 90 - angle1
+        #     width1, length1 = length1, width1
+
+        # Matching longer sides is more important, and thus needs a stricter threshold
+        # if width1 > length1:
+        #     width_sim_thresh, length_sim_thresh = length_sim_thresh, width_sim_thresh
+
+        for j in range(i + 1, len(bboxes)):
+            bbox2 = bboxes[j]
+            center2 = (bbox2[0][0], bbox2[0][1])
+            width2, length2, angle2 = bbox2[1][0], bbox2[1][1], bbox2[2]
+
+            if max(length2, width2) < thresh:
+                continue
+
+            vert2 = cv.boxPoints(bbox2)
+
+            longest = (0, 1)
+            bl = vert2[0]
+            tl = vert2[1]
+            br = vert2[3]
+            width2 = math.sqrt(dist_sq(vert2[0], vert2[3]))
+            if dist_sq(bl, br) > dist_sq(bl, tl):
+                longest = (0, 3)
+                width2 = math.sqrt(dist_sq(vert2[0], vert2[1]))
+            length2 = math.sqrt(dist_sq(vert2[longest[0]], vert2[longest[1]]))
+            angle2 = angle(vert2[longest[0]], vert2[longest[1]])
+
+            # if angle2 > 45:
+            #     angle2 = 90 - angle2
+            #     width2, length2 = length2, width2
+
+            angle_diff = abs(angle1 - angle2)
+            y_diff = abs(center1[1] - center2[1])
+
+            # If two bounding boxes are similar in size and orientation, place a dot between them
+            if width_sim_thresh < sim(width1, width2) \
+                and length_sim_thresh < sim(length1, length2) \
+                and y_diff < y_thresh * (length1 + length2) / 2 \
+                and (angle_diff < angle_thresh or angle_diff > 180 - angle_thresh):
+                cv.circle(frame, (round((bbox1[0][0] + bbox2[0][0]) / 2), round((bbox1[0][1] + bbox2[0][1]) / 2)), 10, (255, 0, 255), -1)
+            cv.putText(frame, f'w: {sim(width1, width2):.2f}, l: {sim(length1, length2):.2f}, a1: {angle1:.2f}, a2: {angle2:.2f}, {angle_diff:.2f}', (round((bbox1[0][0] + bbox2[0][0]) / 2), round((bbox1[0][1] + bbox2[0][1]) / 2)), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
+
+            # Debugging
+            # print(f'w --- {i}: {width1}, {j}: {width2}, sim: {sim(width1, width2)}')
+            # print(f'l --- {i}: {length1}, {j}: {length2}, sim: {sim(length1, length2)}')
+            # print(f'a --- {i}: {angle1}, {j}: {angle2}, diff: {angle_diff}')
+
+    # Reference square to see size of `thresh`
+    # cv.rectangle(frame, (50, 50), (50 + thresh, 50 + thresh), (255, 0, 255), 1)
+
+    return frame
+
+def main():
+    cap = cv.VideoCapture('tests/test1.mp4')
+
+    while cap.isOpened():
+        # start_time = time.time()
+
+        ret, frame = cap.read()
+
+        if not ret:
+            print('Failed to read frame. Exiting...')
+            break
+
+        frame = draw_centers(frame, 'blue')
+
+        cv.imshow('frame', frame)
+
+        # writer.write(frame)
+        # cv.imwrite(output, frame)
+
+        # end_time = time.time()
+        # print(end_time - start_time)
+
+        if cv.waitKey(50) == ord('q'):
+            break
+
+    cap.release()
+    cv.destroyAllWindows()
+
+main()