daily pastebin goal
51%
SHARE
TWEET

Untitled

a guest Feb 22nd, 2019 84 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import numpy as np
  2. import cv2
  3. from math import sqrt
  4.  
  5.  
  6. depth_scale = 0.0010000000474974513
  7. print("Depth Scale is: " , depth_scale)
  8.  
  9. font = cv2.FONT_HERSHEY_SIMPLEX
  10.  
  11.  
  12. color_range =np.array([[20,100,100,25,255,255],[90,250,71,110,255,255],\
  13.                         [80,150,50,90,255,150],[170,190,175,179,255,255],\
  14.                         [12,200,123,19,244,200],[20,00,238,40,255,255],\
  15.                         [0,165,250,8,210,255],[68,120,0,90,255,53]])
  16.  
  17.  
  18. active = []
  19. tbc = []
  20.  
  21. color_c = []
  22. c_h = []
  23.  
  24. try:
  25.     while True:
  26.         # Get frameset of color and depth
  27.         frames = pipeline.wait_for_frames()
  28.         # Align the depth frame to color frame
  29.         aligned_frames = align.process(frames)
  30.         # Get aligned frames
  31.         aligned_depth_frame = aligned_frames.get_depth_frame()
  32.         color_frame = aligned_frames.get_color_frame()
  33.         # Validate that both frames are valid
  34.         if not aligned_depth_frame or not color_frame:
  35.             continue
  36.         # the depth_image array can be used in order to compute the distance,
  37.         # if the value is multiplied with the the depth_scale value
  38.         depth_image = np.asanyarray(aligned_depth_frame.get_data())
  39.         color_image = np.asanyarray(color_frame.get_data())
  40.         # color_image = cv2.imread('color.jpeg',1)
  41.         # np.savetxt('distance.txt', depth_image , fmt='%f')
  42.         # cv2.imwrite('color.jpeg' ,color_image)
  43.         gray = cv2.cvtColor(color_image, cv2.COLOR_BGR2GRAY)
  44.  
  45.  
  46.         gray = cv2.medianBlur(gray,7)
  47.         img = color_image.copy()
  48.         edges = cv2.Canny(gray,100,25) #TODO parameters to be tuned
  49.         c_h *= 0
  50.         circles = cv2.HoughCircles(gray,cv2.HOUGH_GRADIENT,1,30,
  51.                                     param1=50,param2=25,minRadius=10,maxRadius=50)
  52.         if(circles is not None):
  53.             circles = np.reshape(circles, (np.shape(circles)[1],3))
  54.             for i in range(0,np.shape(circles)[0]):
  55.                 cX = circles[i,0]
  56.                 cY = circles[i,1]
  57.                 r = circles[i,2]
  58.                 c_h.append((cX,cY,r))
  59.  
  60.         color_c *= 0
  61.         hsv = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
  62.         thr = np.zeros((8,720,1280))
  63.         # Quite tedious solution, but each color seems to need its own treatment
  64.         for i in range(0,np.shape(color_range)[0]):
  65.             lower = color_range[i,0:3]
  66.             upper = color_range[i,3:6]
  67.             mask = cv2.inRange(hsv, lower, upper)
  68.             if (i == 3):
  69.                 # print('y')
  70.                 mask = cv2.dilate(mask, None, iterations=6)
  71.                 mask2 = cv2.inRange(hsv,(2,213,60),(13,255,180))
  72.                 # mask2 = cv2.dilate(mask2, None, iterations=1)
  73.                 mask = mask2 | mask
  74.                 mask = cv2.dilate(mask, None, iterations=1)
  75.                 mask = cv2.erode(mask, None, iterations=2)
  76.                 mask = cv2.dilate(mask, None, iterations=2)
  77.  
  78.             elif (i==6):
  79.                 mask2 = cv2.inRange(hsv,(176,129,230),(179,156,255))
  80.                 mask = cv2.dilate(mask, None, iterations=7)
  81.                 mask = cv2.erode(mask, None, iterations=2)
  82.                 mask = cv2.dilate(mask, None, iterations=3)
  83.                 mask = mask2 | mask
  84.                 mask = cv2.erode(mask, None, iterations=1)
  85.                 mask = cv2.dilate(mask, None, iterations=1)
  86.             elif (i==7):
  87.                 mask = cv2.erode(mask, None, iterations=2)
  88.                 mask = cv2.dilate(mask, None, iterations=6)
  89.             elif (i==0):
  90.                 mask = cv2.erode(mask, None, iterations=1)
  91.                 mask = cv2.dilate(mask, None, iterations=4)
  92.             elif (i==1):
  93.                 mask = cv2.erode(mask, None, iterations=1)
  94.                 mask = cv2.dilate(mask, None, iterations=4)
  95.             elif (i==2):
  96.                 mask = cv2.erode(mask, None, iterations=2)
  97.                 mask = cv2.dilate(mask, None, iterations=5)
  98.             elif (i==4):
  99.                 mask = cv2.erode(mask, None, iterations=0)
  100.                 mask = cv2.dilate(mask, None, iterations=7)
  101.             elif (i==5):
  102.                 mask = cv2.erode(mask, None, iterations=1)
  103.                 mask = cv2.dilate(mask, None, iterations=6)
  104.             thr[i,:,:] = cv2.threshold(mask,150,255,cv2.THRESH_BINARY)[1]
  105.             cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[1]
  106.             # print('start')
  107.  
  108.             if(len(cnts))> 0:
  109.                 for c in cnts:
  110.                     a = cv2.contourArea(c)
  111.                     if (a > 0 and a < 7850):
  112.                         ((x,y), radius) = cv2.minEnclosingCircle(c)
  113.                         M = cv2.moments(c)
  114.                         if M["m00"] != 0:
  115.                             center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
  116.                         else:
  117.                             center = (0, 0)
  118.                         cX = center[0]
  119.                         cY = center[1]
  120.                         r = radius
  121.  
  122.                         color_c.append((cX,cY,r))
  123.  
  124.         c_final = []
  125.         l_m = np.asarray(color_c)
  126.  
  127.         if (len(active) == 0) and (len(tbc) == 0):
  128.             # this first case basically consist with the initial situation
  129.             # if both active and tbc are emptied, all the new detections have to go into
  130.             # tbc
  131.  
  132.             # the dection is done with a nested for loop, every element is compared to every other.
  133.             # if there is a match within a certain threshold, then there is merge
  134.             for i in range(0,np.shape(l_m)[0]):
  135.                 if (l_m[i][0] >= 0):
  136.                     sum_x = l_m[i][0]
  137.                     sum_y = l_m[i][1]
  138.                     sum_r = l_m[i][2]
  139.                     # r = l_m[i][2]
  140.                     n = 1
  141.                     for j in range(0, np.shape(l_m)[0]):
  142.                         # I don't want to compare the element to itself, nor to an
  143.                         # element which has already been assigned.
  144.                         if(i != j and l_m[j][0]>= 0):
  145.                             r = max(l_m[i][2],l_m[j][2])
  146.                             dis1 = depth_image[int(l_m[i][1]),int(l_m[i][0])]*depth_scale
  147.                             dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
  148.                             if(abs(dis1-dis2)<1.2*r*depth_scale):
  149.                                 if sqrt((l_m[i][0]-l_m[j][0])**2+(l_m[i][1]-l_m[j][1])**2) <= 1.2*r:
  150.                                     # dis1 = depth_image[int(l_m[i][1]),int(l_m[i][0])]*depth_scale
  151.                                     # dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
  152.                                     # if (abs(dis1-dis2)<r):
  153.                                     n = n + 1
  154.                                     sum_x = sum_x + l_m[j][0]
  155.                                     sum_y = sum_y + l_m[j][1]
  156.                                     sum_r = sum_r + l_m[j][2]
  157.                                     # print(type(l_m))
  158.                                     l_m[j][0] = -1
  159.                                     l_m[j][1] = -1
  160.                     cX = int(round(sum_x/n))
  161.                     cY = int(round(sum_y/n))
  162.                     rf = int(round(sum_r/n))
  163.                     tbc.append((cX,cY,rf))
  164.         elif(len(active) == 0) and (len(tbc) != 0):
  165.             # here now tbc has to be compared with the incoming measurements
  166.             t_b = np.asarray(tbc)
  167.             # empty the list, it will be repopulated anyways
  168.             tbc *= 0
  169.             for i in range(0,np.shape(t_b)[0]):
  170.                 flag = 0
  171.                 sum_x = t_b[i][0]
  172.                 sum_y = t_b[i][1]
  173.                 sum_r = t_b[i][2]
  174.                 # r = t_b[i][2]
  175.                 n = 1
  176.                 for j in range(0,np.shape(l_m)[0]):
  177.                     # yeah I cannot have multiple overlaps and like this I won't even be doing anything if a node was already considered before
  178.                     if (l_m[j][0]>0):
  179.                         r = max(t_b[i][2],l_m[j][2])
  180.                         dis1 = depth_image[int(t_b[i][1]),int(t_b[i][0])]*depth_scale
  181.                         dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
  182.                         if(abs(dis1-dis2)<1.2*r*depth_scale):
  183.                             if sqrt((t_b[i][0]-l_m[j][0])**2+(t_b[i][1]-l_m[j][1])**2) <= 1.2*r:
  184.                                 # dis1 = depth_image[int(t_b[i][1]),int(t_b[i][0])]*depth_scale
  185.                                 # dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
  186.                                 # if (abs(dis1-dis2)<r):
  187.                                 flag = 1
  188.                                 sum_x += l_m[j][0]
  189.                                 sum_y += l_m[j][1]
  190.                                 sum_r += l_m[j][2]
  191.                                 n += 1
  192.                                 l_m[j][0] = -1
  193.                                 l_m[j][1] = -1
  194.                 if(flag):
  195.                     # only I have something matches I add it to the list of activte nodes
  196.                     cX = int(round(sum_x/n))
  197.                     cY = int(round(sum_y/n))
  198.                     rf = int(round(sum_r/n))
  199.                     active.append((cX,cY,rf))
  200.             # I don't really like this part
  201.             for k in range(0,np.shape(l_m)[0]):
  202.                 if(l_m[k][0] > 0):
  203.                     tbc.append((l_m[k][0],l_m[k][1],l_m[k][2]))
  204.  
  205.         elif(len(active)!=0):
  206.         # else:
  207.             # this is the last case, in which we should have both tbc, active and the new incoming ones
  208.             t_b = np.asarray(tbc)
  209.             a = np.asarray(active)
  210.             active *= 0
  211.             tbc *= 0
  212.             # I basically do what I did before. Fist the incoming measurments are compared with the active ones. If there are matches the active node is updated
  213.             # otherwise the active nodes will be discarded
  214.             for i in range(0,np.shape(a)[0]):
  215.                 flag = 0
  216.                 sum_x = a[i][0]
  217.                 sum_y = a[i][1]
  218.                 sum_r = a[i][2]
  219.                 r = a[i][2]
  220.                 n = 1
  221.                 for j in range(0,np.shape(l_m)[0]):
  222.                     # yeah I cannot have multiple overlaps and like this I won't even be doing anything if a node was already considered before
  223.                     if (l_m[j][0]>0):
  224.                         r = max(a[i][2],l_m[j][2])
  225.                         dis1 = depth_image[int(a[i][1]),int(a[i][0])]*depth_scale
  226.                         dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
  227.                         if(abs(dis1-dis2)<1.2*r*depth_scale):
  228.                             if sqrt((a[i][0]-l_m[j][0])**2+(a[i][1]-l_m[j][1])**2) <= 1.2*r:
  229.                             #     dis1 = depth_image[int(a[i][1]),int(a[i][0])]*depth_scale
  230.                             #     dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
  231.                                 # if (abs(dis1-dis2)<r):
  232.                                 flag = 1
  233.                                 sum_x += l_m[j][0]
  234.                                 sum_y += l_m[j][1]
  235.                                 sum_r += l_m[j][2]
  236.                                 n += 1
  237.                                 l_m[j][0] = -1
  238.                                 l_m[j][1] = -1
  239.                 if(flag):
  240.                     # only I have something matches I add it to the list of activte nodes
  241.                     cX = int(round(sum_x/n))
  242.                     cY = int(round(sum_y/n))
  243.                     rf = int(round(sum_r/n))
  244.                     active.append((cX,cY,rf))
  245.             # After comparing to the active ones, I do the comparison to the tbc centers, If there are any matches
  246.             # I add to active. If there is no match then I simply add the new detection to tbc
  247.             for i in range(0,np.shape(t_b)[0]):
  248.                 flag = 0
  249.                 sum_x = t_b[i][0]
  250.                 sum_y = t_b[i][1]
  251.                 sum_r = t_b[i][2]
  252.                 r = t_b[i][2]
  253.                 n = 1
  254.                 for j in range(0,np.shape(l_m)[0]):
  255.                     # yeah I cannot have multiple overlaps and like this I won't even be doing anything if a node was already considered before
  256.                     if (l_m[j][0]>0):
  257.                         r = max(t_b[i][2],l_m[j][2])
  258.                         dis1 = depth_image[int(t_b[i][1]),int(t_b[i][0])]*depth_scale
  259.                         dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
  260.                         if(abs(dis1-dis2)<1.2*r*depth_scale):
  261.                             if sqrt((t_b[i][0]-l_m[j][0])**2+(t_b[i][1]-l_m[j][1])**2) <= 1.2*r:
  262.                                 # dis1 = depth_image[int(t_b[i][1]),int(t_b[i][0])]*depth_scale
  263.                                 # dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
  264.                                 # if (abs(dis1-dis2)<r):
  265.                                 flag = 1
  266.                                 sum_x += l_m[j][0]
  267.                                 sum_y += l_m[j][1]
  268.                                 sum_r += l_m[j][2]
  269.                                 n += 1
  270.                                 l_m[j][0] = -1
  271.                                 l_m[j][1] = -1
  272.                 if(flag):
  273.                     # only I have something matches I add it to the list of activte nodes
  274.                     cX = int(round(sum_x/n))
  275.                     cY = int(round(sum_y/n))
  276.                     rf = int(round(sum_r/n))
  277.                     active.append((cX,cY,rf))
  278.             for k in range(0,np.shape(l_m)[0]):
  279.                 if(l_m[k][0] > 0):
  280.                     tbc.append((l_m[k][0],l_m[k][1],l_m[k][2]))
  281.         red = []
  282.         yellow = []
  283.         pink = []
  284.         black = []
  285.         white = []
  286.         blue = []
  287.         green = []
  288.         brown = []
  289.         # this is what is needed in order to check whether we are dealing with a ball which is on the table
  290.         lower = np.asarray([60,100,100])
  291.         upper = np.asarray([80,255,255])
  292.         mask = cv2.inRange(hsv, lower, upper)
  293.         border = cv2.threshold(mask,150,255,cv2.THRESH_BINARY)[1]
  294.  
  295.         for z in range(0,np.shape(active)[0]):
  296.                 cX, cY, r = active[z]
  297.                 dis = depth_image[int(cY)-1,int(cX)]*depth_scale
  298.                 dis = round(dis,4)
  299.                 # check whether a minimum distance is fulfilled
  300.                 if(dis < 5.2 and dis> 0.2):
  301.                     r_new = int(round(r*2))
  302.                     cX_ma = cX + r_new
  303.                     cX_mi = cX - r_new
  304.                     cY_ma = cY + r_new
  305.                     cY_mi = cY - r_new
  306.                     if(cX_ma > 1279):
  307.                         cX_ma = 1279
  308.                     if(cX_mi < 0):
  309.                         cX_mi = 0
  310.                     if(cY_ma > 719):
  311.                         cY_ma = 719
  312.                     if(cY_mi < 0):
  313.                         cY_mi = 0
  314.                     # check if the ball is on the table
  315.                     if(((border[cY_mi,cX]!=0)) or ((border[cY_ma,cX]!=0)) or ((border[cY,cX_mi]!=0)) or ((border[cY,cX_ma])!=0) ):
  316.                       if np.sum(thr[0,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
  317.                                   (int(round(active[z][0]-15))):int(round((active[z][0]+10)))])/255 > 40 :
  318.                           yellow.append((active[z]))
  319.                       elif np.sum(thr[2,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
  320.                                   (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
  321.                           green.append((active[z]))
  322.                       elif np.sum(thr[1,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
  323.                                   (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
  324.                           blue.append((active[z]))
  325.                       elif np.sum(thr[3,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
  326.                                   (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
  327.                           red.append((active[z]))
  328.                       elif( np.sum(thr[6,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
  329.                               int(round((active[z][0]-15))):int(round((active[z][0]+15)))])/255 >40):
  330.                           pink.append((active[z]))
  331.                       elif np.sum(thr[4,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
  332.                                     (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
  333.                           brown.append((active[z]))
  334.                       elif np.sum(thr[5,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
  335.                                   (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
  336.                           white.append((active[z]))
  337.                       elif np.sum(thr[7,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
  338.                                   (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
  339.                           black.append((active[z]))
  340.         all = red + yellow + blue + green +brown + white + pink + black
  341.         active *= 0
  342.         active = list(all)
  343.         # print(brown)
  344.         t = np.asarray(all)
  345.         # t = np.asarray(red)
  346.  
  347.         # print(t)
  348.         # t = np.asarray(brown)
  349.         # print('white')
  350.         # print(white)
  351.         # print('blue')
  352.         # print(blue)
  353.  
  354.  
  355.         # cv2.imshow('m', mask)
  356.  
  357.         # print(border[100,200])
  358.         # print('start')
  359.         for n in range(0, np.shape(t)[0]):
  360.             cX, cY, r = t[n]
  361.             dis = depth_image[int(cY),int(cX)]*depth_scale
  362.             dis = round(dis,4)
  363.             r_new = int(round(r*1.5))
  364.             cX_ma = cX + r_new
  365.             cX_mi = cX - r_new
  366.             cY_ma = cY + r_new
  367.             cY_mi = cY - r_new
  368.             cv2.rectangle(color_image,(cX-r_new,cY-r_new),(cX+r_new,cY+r_new),(0, 255, 0), 1)
  369.             cv2.putText(color_image, "Distance: " + str(dis), (int(cX)-r_new,int(cY)+2*r_new), font, 1, (0,0,255),2,cv2.LINE_AA)
  370.  
  371.  
  372.         depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET)
  373.         cv2.imwrite("distance_image.png",color_image)
  374.         # cv2.imwrite("dept_image.png", depth_colormap)
  375.         images = np.hstack((color_image, depth_colormap))
  376.         cv2.namedWindow('Align Example', cv2.WINDOW_AUTOSIZE)
  377.         cv2.imshow('Align Example', images)
  378.         # cv2.imshow('red2', mask2)
  379.         cv2.imshow('thr',thr[7,:,:])
  380.         cv2.imshow('b',border)
  381.         # cv2.imshow('Gray',gray)
  382.         key = cv2.waitKey(1)
  383.         # Press esc or 'q' to close the image window
  384.         if key & 0xFF == ord('q') or key == 27:
  385.             cv2.destroyAllWindows()
  386.             break
  387. finally:
  388.     pipeline.stop()
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