Merge pull request #3 from Blenders-FC/package_structure

[PKG] Package structure

Author: pedrodeniz48

launch/vision.launch

@@ -0,0 +1,14 @@
+<?xml version="1.0" ?>
+<launch>    
+
+  <!-- Define robot_id with a default value (used if not inherited) -->
+  <arg name="robot_id" default="1"/>
+  <param name="robot_id" value="$(arg robot_id)"/>
+
+  <!-- Camera Node -->
+  <node pkg="usb_cam" type="usb_cam_node" name="usb_cam_node" output="screen"/>
+
+  <!-- Vision Node -->
+  <node pkg="vision_pkg" type="main_vision_node.py" name="main_vision_node" output="screen"/>
+
+</launch>

src/_classes.txt models/_classes.txtsrc/_classes.txt renamed to models/_classes.txt

@@ -21,8 +21,21 @@ def calculate_fps(prev_frame_time, new_frame_time, frame, position=(7, 70), font
     cv2.putText(frame, fps_text, position, font, font_scale, color, thickness, cv2.LINE_AA)
     return prev_frame_time
 
+# Function to calculate lines angles
+def calculate_angles(lines):
+    try:
+        for line in lines:
+            x1, y1, x2, y2 = line[0]
+            # Calculate the angle of the line
+            angle = np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi  # Convert to degrees
+            print(f"Line angle: {angle} degrees")
+    except TypeError:
+        pass
+
+
+
 # Open the video file
-video_path = Path(__file__).parent.parent / "media" / "video1.mp4"  # Get path relative to script location
+video_path = Path(__file__).resolve().parent.parent / "media" / "video1.mp4"  # Get path relative to script location
 cap = cv2.VideoCapture(str(video_path))
 
 # Check if the video opened successfully
@@ -38,12 +51,18 @@ def calculate_fps(prev_frame_time, new_frame_time, frame, position=(7, 70), font
 # used to record the time when we processed last frame 
 prev_frame_time = 0
 
-# Pre-calculate zones to exclude from drawing lines
-def calculate_nozone(shape, offset=50):
-    y, x = shape[:2]
-    y_nozone = list(range(y - offset, y + 1)) + list(range(0, offset + 1))
-    x_nozone = list(range(x - offset, x + 1)) + list(range(0, offset + 1))
-    return y_nozone, x_nozone
+# Calculate zones for the current frame
+ret, orig_frame = cap.read()
+
+# Define the region of interest (ROI)
+image_height, image_width = orig_frame.shape[:2]
+roi_height = int(image_height * 0.5)  # Crop the top 50% of the image
+roi_width = int(image_width * 0.8)    # Crop the middle 80% of the image
+x_offset = int(image_width * 0.1)     # Offset to focus on the center
+y_offset = int(image_height * 0.25)  # Vertical offset
+
+orig_frame = orig_frame[y_offset:y_offset+roi_height, x_offset:x_offset+roi_width]
+yimg_nozone, ximg_nozone = calculate_nozone(orig_frame.shape)
 
 # Loop to read frames
 while True:
@@ -52,12 +71,12 @@ def calculate_nozone(shape, offset=50):
     if not ret:
         break
 
+    # Crop the image
+    orig_frame = orig_frame[y_offset:y_offset+roi_height, x_offset:x_offset+roi_width]
+
     # Start measuring time for each frame
     new_frame_time = time.time()
 
-    # Calculate zones for the current frame
-    yimg_nozone, ximg_nozone = calculate_nozone(orig_frame.shape)
-    
     # Blur and convert to grayscale
     blurred = cv2.GaussianBlur(orig_frame, (9, 9), 0)
     gray = cv2.cvtColor(blurred, cv2.COLOR_BGR2GRAY)
@@ -85,6 +104,7 @@ def calculate_nozone(shape, offset=50):
 
     # Find contours of the field markings
     contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:10]  # Keep top 10 largest contours
 
     # Draw only large contours
     dummy_mask = np.zeros_like(mask_green)
@@ -95,20 +115,27 @@ def calculate_nozone(shape, offset=50):
     #adap_thrs_morph = cv2.morphologyEx(dummy_mask, cv2.MORPH_CLOSE, kernel2, iterations=1)
 
     # Detect lines using Hough transform
-    lines = cv2.HoughLinesP(dummy_mask, rho=2, theta=np.pi / 180, threshold=250,
-                            minLineLength=25, maxLineGap=50)
+    lines = cv2.HoughLinesP(dummy_mask, rho=10, theta=np.pi / 180, threshold=250,
+                            minLineLength=10, maxLineGap=50)  # minLineLength=25
+    
+    filtered_lines = []
+    if lines is not None:
+        for line in lines:
+            x1, y1, x2, y2 = line[0]
+            if (y1 not in yimg_nozone or y2 not in yimg_nozone) and (x1 not in ximg_nozone or x2 not in ximg_nozone):
+                filtered_lines.append(line)
+
+    calculate_angles(filtered_lines)
 
     # Create a copy to draw lines on
     image_lines_thrs = orig_frame.copy()
     lines_thrs = np.zeros_like(dummy_mask)
 
     # Draw detected lines while avoiding no-line zones
-    if lines is not None:
-        for line in lines:
-            x1, y1, x2, y2 = line[0]
-            if (y1 not in yimg_nozone or y2 not in yimg_nozone) and (x1 not in ximg_nozone or x2 not in ximg_nozone):
-                cv2.line(image_lines_thrs, (x1, y1), (x2, y2), (255, 0, 0), 10)  # Blue lines
-                cv2.line(lines_thrs, (x1, y1), (x2, y2), (255, 0, 0), 10)
+    for line in filtered_lines:
+        x1, y1, x2, y2 = line[0]
+        cv2.line(image_lines_thrs, (x1, y1), (x2, y2), (255, 0, 0), 1)  # Blue lines
+        cv2.line(lines_thrs, (x1, y1), (x2, y2), (255, 0, 0), 10)
 
     # Mask lines with white areas
     lines_thrs = cv2.bitwise_and(lines_thrs, white_msk)
@@ -125,6 +152,7 @@ def calculate_nozone(shape, offset=50):
 
     # Display the combined frame
     cv2.imshow('Combined Video', overlay_image)
+    # cv2.imshow("Region of Interest", roi)
 
     # Break loop if 'q' is pressed
     if cv2.waitKey(30) & 0xFF == ord('q'):

src/lines_detector.py

@@ -2,16 +2,17 @@
 
 import cv2
 import numpy as np
-from utils import *
-from openvino.runtime import Core
 import rospy
 import math
+from yolo_utils import *
+from openvino.runtime import Core
 from sensor_msgs.msg import Image
 from sensor_msgs.msg import JointState
 from cv_bridge import CvBridge, CvBridgeError
 from geometry_msgs.msg import Point
 from std_msgs.msg import Float64
 
+
 def detect(session, image_src, namesfile):
     global center
 
@@ -34,15 +35,15 @@ def detect(session, image_src, namesfile):
 
     boxes = post_processing(img=img_in, conf_thresh=0.8, nms_thresh=0.6, output=outputs)
 
-    print(boxes)
+    # print(boxes)
 
     if len(boxes) > 0 and len(boxes[0]) > 0 and len(boxes[0][0]) > 0:
-        print(boxes[0])
+        # print(boxes[0])
         # Update center coordinates
         center.x = int(boxes[0][0][7] * IN_IMAGE_W)
         center.y = int(boxes[0][0][8] * IN_IMAGE_H)
         center.z = int((boxes[0][0][2]*IN_IMAGE_W - boxes[0][0][0]*IN_IMAGE_W ) * (boxes[0][0][3]*IN_IMAGE_H- boxes[0][0][1]*IN_IMAGE_H))
-        print(center)
+        # print(center)
         pub_center.publish(center)
     else:
         center.x = 999
@@ -54,18 +55,19 @@ def detect(session, image_src, namesfile):
 
     return final_img, center
 
+
 def compileModel():
     # Function for compiling model with OpenVino
     global compiled_model
 
     ie = Core()
-    model = ie.read_model(model="/home/robotis/blenders_ws/src/soccer_pkg/src/yolov4soccer.onnx")
-    compiled_model = ie.compile_model(model=model, device_name="MYRIAD")
+    model = ie.read_model(model=find_file(parent_folder="models", filename="yolov4soccer.onnx"))
+    compiled_model = ie.compile_model(model=model, device_name="CPU")  # "MYRIAD")
 
 
 def imageCallback(img_msg):
     global center, compiled_model
-    print("Llega image callback")
+    # print("Llega image callback")
     try:
         frame = bridge.imgmsg_to_cv2(img_msg, "passthrough")
     except CvBridgeError as e:
@@ -74,7 +76,7 @@ def imageCallback(img_msg):
 
     # Process image with model and publish
 
-    namesfile = '/home/robotis/blenders_ws/src/soccer_pkg/src/_classes.txt'
+    namesfile = find_file(parent_folder="models", filename="_classes.txt")
 
     img, center = detect(session=compiled_model, image_src=frame, namesfile=namesfile)
     final_img = bridge.cv2_to_imgmsg(img, "rgb8")
@@ -93,18 +95,15 @@ def imageCallback(img_msg):
 
     compileModel()
 
-    
-
-    rospy.init_node('yolo_vision')
+    rospy.init_node('vision_node')
     robot_id = rospy.get_param('robot_id', 1)
 
 
     pub_img = rospy.Publisher(f'/robotis_{robot_id}/ImgFinal', Image, queue_size=1)
     pub_center = rospy.Publisher(f'/robotis_{robot_id}/ball_center', Point, queue_size=1)
 
-    rospy.loginfo("Hello ROS")
+    rospy.loginfo("Vision node initialized")
 
     subimg = rospy.Subscriber("/usb_cam_node/image_raw", Image, imageCallback)
 
     rospy.spin()
-

src/yolo_node.py src/main_vision_node.pysrc/yolo_node.py renamed to src/main_vision_node.py

@@ -4,6 +4,7 @@
 import numpy as np
 import cv2 as cv
 from collections import deque
+from pathlib import Path
 
 
 def sigmoid(x):
@@ -116,7 +117,7 @@ def get_color(c, x, max_val):
         if len(box) >= 7 and class_names:
             cls_conf = box[5]
             cls_id = box[6]
-            print('%s: %f' % (class_names[cls_id], cls_conf))
+            # print('%s: %f' % (class_names[cls_id], cls_conf))
             classes = len(class_names)
             offset = cls_id * 123457 % classes
             red = get_color(2, offset, classes)
@@ -140,11 +141,12 @@ def get_color(c, x, max_val):
         cv2.imwrite(savename, img)
     return img
 
+
 def plot_boxes_cv2_video(img, boxes, class_names=None, color=None):
     import cv2
     img = np.copy(img)
     colors = np.array([[1, 0, 1], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 0]], dtype=np.float32)
-    print(boxes)
+    # print(boxes)
     def get_color(c, x, max_val):
         ratio = float(x) / max_val * 5
         i = int(math.floor(ratio))
@@ -263,6 +265,19 @@ def post_processing(img, conf_thresh, nms_thresh, output):
     t3 = time.time()
 
     return bboxes_batch
+
+
+def find_file(parent_folder, filename):
+    current_dir = Path(__file__).resolve().parent
+    soccer_pkg_dir = current_dir.parent
+
+    dir_path = soccer_pkg_dir / parent_folder
+    file_path = dir_path / filename
+
+    if file_path.exists():
+        return file_path
+    else:
+        return None
 
 
 class CvFpsCalc(object):