SHARE
TWEET

line_detec

a guest May 25th, 2019 80 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/usr/bin/env python
  2.  
  3. # Python libs
  4. import os
  5. import sys, time
  6. import png
  7. import math
  8. from sklearn.cluster import KMeans
  9.  
  10. from LineDetection import LineDetection
  11. # numpy and scipy
  12. import numpy as np
  13. from scipy.ndimage import filters
  14. from PIL import Image
  15.  
  16. import matplotlib
  17. from std_msgs.msg import Float64
  18.  
  19. # OpenCV
  20. import cv2
  21.  
  22. # Ros libraries
  23. import roslib
  24. import rospy
  25. from math import sqrt
  26.  
  27. # Ros Messages
  28. from sensor_msgs.msg import CompressedImage
  29. # We do not use cv_bridge it does not support CompressedImage in python
  30. # from cv_bridge import CvBridge, CvBridgeError
  31.  
  32. VERBOSE=False
  33. ITERATIONS = 5
  34. DIFF_FACTOR = 10
  35.  
  36. class image_feature:
  37.  
  38.     def __init__(self):
  39.         self.curr_cam = None
  40.  
  41.         # subscribed Topic
  42.         self.subscriber = rospy.Subscriber("husky/camera1/image_raw/compressed",
  43.             CompressedImage, self.callback,  queue_size = 10)
  44.  
  45.         self.image_pub = rospy.Publisher(
  46.             "/output/image_raw/compressed",
  47.             CompressedImage,
  48.             queue_size=1)
  49.  
  50.         self.angle_pub = rospy.Publisher("/detected_angle", Float64, queue_size=1)
  51.  
  52.         if VERBOSE :
  53.             print ("subscribed to /camera/image/compressed")
  54.  
  55.         self.color = [0, 0, 0]
  56.         self.color_initilized = False
  57.         self.kmeans_iter = 0
  58.            
  59.     def process(self):
  60.         if self.curr_cam is None:
  61.             print("Initialize self.curr_cam")
  62.             return
  63.  
  64.         if not self.color_initilized:
  65.             print("Color not initialized, exiting detection")
  66.             self.publish_angle(-1)
  67.             return
  68.  
  69.         #### direct conversion to CV2 ####
  70.         np_arr = np.fromstring(self.curr_cam, np.uint8)
  71.         image_np = cv2.imdecode(np_arr, cv2.COLOR_BGR2RGB)
  72.  
  73.         h = np.size(image_np, 0)
  74.         w = np.size(image_np, 1)
  75.  
  76.         w, h = image_np.shape[:2]
  77.         image_np = image_np[h/2:h, 0:w].copy()
  78.         image_orig = image_np
  79.  
  80.         thresh = 20
  81.         greenLower = (self.color[0] - thresh, self.color[1] - thresh, self.color[2] - thresh)
  82.         greenUpper = (self.color[0] + thresh, self.color[1] + thresh, self.color[2] + thresh)
  83.         print("greenupper: ", greenUpper)
  84.         print("greenlower; ", greenLower)
  85.         print("\n")
  86.  
  87.         image_np = cv2.inRange(image_np, greenLower, greenUpper)
  88.        
  89.         edges = cv2.Canny(image_np, 50, 150, apertureSize=3)
  90.         minLineLength = 5      # Minimum length of line. Line segments shorter than this are rejected.
  91.         maxLineGap = 100          # Maximum allowed gap between line segments to treat them as single line.
  92.         rho_precision = 1
  93.         theta_precision = np.pi / 180
  94.         vote_count = 50
  95.  
  96.         lines = cv2.HoughLines(edges, rho_precision, theta_precision, vote_count) # minLineLength, maxLineGap)
  97.         # If no lines are found punish and continue
  98.         if lines is None:
  99.             print("CameraProcessing.run() - no lines found")
  100.             avg_theta = 500
  101.         try:
  102.             avg, img = self.draw_hough_lines(lines, image_np, image_orig)
  103.             self.publish_angle(avg)
  104.  
  105.         except:
  106.             print("There is no detected line.")
  107.  
  108.             self.publish_img(image_np, self.image_pub)
  109.             self.publish_angle(-1)
  110.  
  111.             return
  112.  
  113.         self.publish_img(img, self.image_pub)
  114.         print(avg)
  115.         #self.subscriber.unregister()
  116.  
  117.      def initialize_color(self):
  118.        
  119.         if self.curr_cam is None:
  120.             print("Initialize self.curr_cam")
  121.             return
  122.  
  123.         if self.color_initilized:
  124.             return
  125.  
  126.         np_arr = np.fromstring(self.curr_cam, np.uint8)
  127.         image_np = cv2.imdecode(np_arr, cv2.COLOR_BGR2RGB)
  128.  
  129.         # Do the kmeans
  130.         image = image_np.reshape((image_np.shape[0] * image_np.shape[1], 3))
  131.         clt = KMeans(n_clusters = 5)
  132.         clt.fit(image)
  133.         print("Center clusters:", clt.cluster_centers_)
  134.        
  135.         diff = 0
  136.         ind = -1
  137.         # find the most different cluster
  138.         for i, c in enumerate(clt.cluster_centers_):
  139.             d = sqrt( ( c[0] - c[1]) ** 2 + (c[1] - c[2])**2 )
  140.             if d > diff:
  141.                 diff = d
  142.                 ind = i
  143.  
  144.         if diff < DIFF_FACTOR:
  145.             print("unable to find a different color")
  146.         else:
  147.             self.kmeans_iter += 1
  148.             self.color[0] += clt.cluster_centers_[ind][0]
  149.             self.color[1] += clt.cluster_centers_[ind][1]
  150.             self.color[2] += clt.cluster_centers_[ind][2]
  151.             print("Found diff color:", clt.cluster_centers_[ind])
  152.             print("Iteration count {}", self.kmeans_iter)
  153.  
  154.         if self.kmeans_iter == ITERATIONS:
  155.             self.color_initilized = True
  156.             self.color[0] /= ITERATIONS
  157.             self.color[1] /= ITERATIONS
  158.             self.color[2] /= ITERATIONS
  159.             print("Using color: ", self.color)
  160.  
  161.     def callback(self, ros_data):
  162.         self.curr_cam = ros_data.data
  163.        
  164.  
  165.     def publish_img(self, img, pub):
  166.         # Create published image
  167.         msg = CompressedImage()
  168.         msg.header.stamp = rospy.Time.now()
  169.         msg.format = "jpeg"
  170.         msg.data = np.array(cv2.imencode('.jpg', img)[1]).tostring()
  171.         pub.publish(msg)
  172.  
  173.     def publish_angle(self, angle):
  174.         # Create message
  175.         msg = Float64()
  176.         msg.data = angle
  177.         self.angle_pub.publish(msg)
  178.  
  179.     def draw_hough_lines(self, lines, img, orig_img):
  180.         """
  181.        Draw Hough lines on the given image
  182.  
  183.        :param lines: Line array.
  184.        :param img: Given image.
  185.  
  186.        :return: Image with drawn lines
  187.        """
  188.  
  189.         avg_theta = 0
  190.         draw_theta = 0
  191.         draw_rho = 0
  192.         for line in lines:
  193.             # Extract line info
  194.  
  195.             rho = line[0][0]
  196.             theta = line[0][1]
  197.  
  198.             draw_theta += theta
  199.             draw_rho += rho
  200.  
  201.             if theta > math.pi:
  202.                 theta = theta - math.pi
  203.  
  204.             theta = math.pi/2-theta
  205.  
  206.             if theta < 0:
  207.                 theta = theta + math.pi
  208.  
  209.             avg_theta += theta
  210.  
  211.         avg_theta /= len(lines)
  212.         draw_theta /= len(lines)
  213.         draw_rho /= len(lines)
  214.  
  215.         a = np.cos(draw_theta)
  216.         b = np.sin(draw_theta)
  217.         x0 = a*draw_rho
  218.         y0 = b*draw_rho
  219.         x1 = int(x0 + 2000*(-b))
  220.         y1 = int(y0 + 2000*(a))
  221.         x2 = int(x0 - 2000*(-b))
  222.         y2 = int(y0 - 2000*(a))
  223.  
  224.         cv2.line(orig_img, (x1, y1), (x2, y2), (0, 0, 255), 5)
  225.  
  226.         final_angle = avg_theta * 180 / math.pi
  227.  
  228.         return final_angle, orig_img
  229.  
  230. def main(args):
  231.     '''Initializes and cleanup ros node'''
  232.     rospy.init_node("line_detection_node")
  233.     ic = image_feature()
  234.    
  235.     while not rospy.is_shutdown():
  236.         ic.initialize_color()
  237.         ic.process()
  238.         rospy.sleep(0.05)
  239.  
  240.     cv2.destroyAllWindows()
  241.  
  242. if __name__ == '__main__':
  243.     main(sys.argv)
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top