Advertisement
patryk

KCK - FindGuitars

Nov 22nd, 2015
393
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 8.83 KB | None | 0 0
  1. #!/usr/bin/python
  2. #-*coding: utf-8 -*-
  3.  
  4. import numpy as np
  5. import cv2
  6. import sys
  7. import matplotlib.pyplot as plt
  8.  
  9.  
  10. '''
  11.     PRZETWARZANIE OBRAZU -
  12.     WYKRYWANIE I ROZRÓŻNIANIE GITARY ELEKTRYCZNEJ OD AKUSTYCZNEJ
  13.  
  14.     Autorzy: Katarzyna Boczek, Patryk Gliszczyński
  15.  
  16.  
  17.     Program ten ma na celu rozpoznawanie dwóch charakterystycznych
  18.     typów gitar na podstawie otrzymanego obrazu i zaznaczenie ich
  19.     położenia. Ponieważ gitara w odróżnieniu od na przykład tablic
  20.     rejestracyjnych czy kart do gry ma bardzo nieregularny kształt,
  21.     przez co przy różnych ustawieniach gitary bardzo utrudnia weryfikacje
  22.     jej pozycji. Dlatego też mimo wielu prób nie potrafiliśmy napisać,
  23.     jednej metody, która by w każdych warunkach bezbłędnie wykrywała
  24.     położenie gitary. Oczywiście w idealnych warunkach, gdy gitara jest
  25.     do nas skierowana bezpośrednio nie było żadnego problemu w wykryciu
  26.     gitary na obrazie. Natomiast zbyt mocna zmiana perspektywy i położenie
  27.     gitary było prawie niemożliwe do detekcji. Próbowaliśmy porównywać z
  28.     wieloma "zdjęciami podstawowymi" np. sam gryf, sama podstrunnica, czy
  29.     też gitara w całości, natomiast nie znaleźliśmy złotego środka, które
  30.     zdjęcie byłoby najlepsze do porównania go z otrzymanym obrazem.
  31. '''
  32.  
  33.  
  34.  
  35. #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  36. # FIND GUITARS
  37.  
  38. def findGuitars(fn1, fn2, fn3, method):
  39.     electricImg = cv2.imread(fn1, 0)
  40.     acousticImg = cv2.imread(fn2, 0)
  41.     mainImg = cv2.imread(fn3, 0)
  42.     electricColorImg = cv2.imread(fn1)
  43.     acousticColorImg = cv2.imread(fn2)    
  44.     mainColorImg = cv2.imread(fn3)
  45.  
  46.     if electricImg is None:
  47.         print ('Failed to load fn1:', fn1)
  48.         sys.exit(1)
  49.        
  50.     if acousticImg is None:
  51.         print ('Failed to load fn2:', fn2)
  52.         sys.exit(1)
  53.  
  54.     if mainImg is None:
  55.         print ('Failed to load fn3:', fn3)
  56.         sys.exit(1)
  57.  
  58.     kp_pairs1 = match_images(electricImg, mainImg, method)
  59.     kp_pairs2 = match_images(acousticImg, mainImg, method)
  60.    
  61.     if kp_pairs1:
  62.         draw_matches('find_obj', kp_pairs1, kp_pairs2, electricImg, acousticImg,
  63.         mainImg, electricColorImg, acousticColorImg, mainColorImg)  
  64.     else:
  65.         print ("No matches found")
  66.        
  67.  
  68. def match_images(img1, img2, method):
  69.  
  70.     if (method == 'SIFT'):
  71.         detector = cv2.SIFT()
  72.     elif (method == 'SURF'):
  73.         detector = cv2.SURF(300, 5, 5)
  74.     elif (method == 'ORB'):
  75.         detector = cv2.ORB()
  76.     elif (method == 'BRIEF'):
  77.         star = cv2.FeatureDetector_create("STAR")
  78.         brief = cv2.DescriptorExtractor_create("BRIEF")
  79.         matcher = cv2.BFMatcher(cv2.NORM_L2)
  80.         kp1 = star.detect(img1, None)
  81.         kp2 = star.detect(img2, None)
  82.         kp1, desc1 = brief.compute(img1, kp1)
  83.         kp2, desc2 = brief.compute(img2, kp2)
  84.  
  85.         raw_matches = matcher.knnMatch(desc1, trainDescriptors = desc2, k = 2)
  86.         kp_pairs = filter_matches(kp1, kp2, raw_matches)
  87.         return kp_pairs
  88.    
  89.     else:
  90.         print ("Wrong Method!")
  91.         sys.exit(0)
  92.  
  93.     matcher = cv2.BFMatcher(cv2.NORM_L2)
  94.  
  95.     kp1, desc1 = detector.detectAndCompute(img1, None)
  96.     kp2, desc2 = detector.detectAndCompute(img2, None)
  97.  
  98.     raw_matches = matcher.knnMatch(desc1, trainDescriptors = desc2, k = 2)
  99.     kp_pairs = filter_matches(kp1, kp2, raw_matches)
  100.     return kp_pairs
  101.  
  102.  
  103. def filter_matches(kp1, kp2, matches, ratio = 0.75):
  104.     mkp1, mkp2 = [], []
  105.    
  106.     for m in matches:
  107.         if len(m) == 2 and m[0].distance < m[1].distance * ratio:
  108.             m = m[0]
  109.             mkp1.append( kp1[m.queryIdx] )
  110.             mkp2.append( kp2[m.trainIdx] )
  111.    
  112.     kp_pairs = zip(mkp1, mkp2)
  113.     return kp_pairs
  114.    
  115.  
  116. def explore_match(win, kp_pairs1, kp_pairs2, electricImg,
  117.         acousticImg, mainImg, electricColorImg, acousticColorImg,
  118.         mainColorImg, status1=None, H1=None, status2=None, H2=None):
  119.    
  120.     h1, w1 = electricImg.shape[:2]
  121.     h2, w2 = mainImg.shape[:2]
  122.     h3, w3 = acousticImg.shape[:2]
  123.    
  124.     vis = np.zeros((max(h1, h2, h3), w1+w2+w3), np.uint8)
  125.     vis = cv2.cvtColor(vis, cv2.COLOR_GRAY2BGR)
  126.     vis[:h1, :w1] = electricColorImg
  127.     vis[:h2, w1:w1+w2] = mainColorImg
  128.     vis[:h3, w1+w2:w1+w2+w3] = acousticColorImg
  129.  
  130.     green = (0, 255, 0)
  131.     red = (0, 0, 255)
  132.     electricColor = (10, 10, 255)
  133.     acousticColor = (0,255,255)
  134.     electricLines = (10,255,10)
  135.     acousticLines = (10,255,10)
  136.    
  137.     if H1 is not None:
  138.         corners1 = np.float32([[0, 0], [w1, 0], [w1, h1], [0, h1]])
  139.         corners1 = np.int32( cv2.perspectiveTransform(corners1.reshape(1, -1, 2), H1).
  140.             reshape(-1, 2) + (w1, 0) )
  141.         cv2.polylines(vis, [corners1], True, electricColor, 3)
  142.  
  143.     if status1 is None:
  144.         status1 = np.ones(len(kp_pairs1), np.bool_)
  145.    
  146.     p1_1 = np.int32([kpp[0].pt for kpp in kp_pairs1])
  147.     p2_1 = np.int32([kpp[1].pt for kpp in kp_pairs1]) + (w1, 0)
  148.  
  149.  
  150.     if H2 is not None:
  151.         corners2 = np.float32([[0, 0], [w3, 0], [w3, h3], [0, h3]])
  152.         corners2 = np.int32( cv2.perspectiveTransform(corners2.reshape(1, -1, 2), H2).
  153.             reshape(-1, 2) + (w1, 0))
  154.         cv2.polylines(vis, [corners2], True, acousticColor, 3)
  155.  
  156.     if status2 is None:
  157.         status2 = np.ones(len(kp_pairs2), np.bool_)
  158.     p1_2 = np.int32([kpp[0].pt for kpp in kp_pairs2])
  159.     p2_2 = np.int32([kpp[1].pt for kpp in kp_pairs2]) + (w3, 0)
  160.  
  161.    
  162.     for (x1, y1), (x2, y2), inlier in zip(p1_1, p2_1, status1):
  163.         if inlier:
  164.             col = green
  165.             thickness = 3
  166.             cv2.circle(vis, (x1, y1), 2, col, -1)
  167.             cv2.circle(vis, (x2, y2), 2, col, -1)
  168.         else:
  169.             col = red
  170.             r = 2
  171.             thickness = 3
  172.             cv2.line(vis, (x1-r, y1-r), (x1+r, y1+r), col, thickness)
  173.             cv2.line(vis, (x1-r, y1+r), (x1+r, y1-r), col, thickness)
  174.             cv2.line(vis, (x2-r, y2-r), (x2+r, y2+r), col, thickness)
  175.             cv2.line(vis, (x2-r, y2+r), (x2+r, y2-r), col, thickness)
  176.  
  177.     for (x1, y1), (x2, y2), inlier in zip(p1_2, p2_2, status2):
  178.         if inlier:
  179.             col = green
  180.             thickness = 3
  181.             cv2.circle(vis, (x1 + w1 + w2 , y1), 2, col, -1)
  182.             cv2.circle(vis, (x2 + w1 - w3, y2), 2, col, -1)
  183.         else:
  184.             col = red
  185.             r = 2
  186.             thickness = 3
  187.             cv2.line(vis, (x1-r+(w1+w2), y1-r), (x1+r+(w1+w2), y1+r), col, thickness)
  188.             cv2.line(vis, (x1-r+(w1+w2), y1+r), (x1+r+(w1+w2), y1-r), col, thickness)
  189.             cv2.line(vis, (x2-r+(w1-w3), y2-r), (x2+r+(w1-w3), y2+r), col, thickness)
  190.             cv2.line(vis, (x2-r+(w1-w3), y2+r), (x2+r+(w1-w3), y2-r), col, thickness)
  191.    
  192.     vis0 = vis.copy()
  193.    
  194.     if (H1 is not None):
  195.         for (x1, y1), (x2, y2), inlier in zip(p1_1, p2_1, status1):
  196.             if inlier:
  197.                 cv2.line(vis, (x1, y1), (x2, y2), electricLines)
  198.  
  199.     if (H2 is not None):
  200.         for (x1, y1), (x2, y2), inlier in zip(p1_2, p2_2, status2):
  201.             if inlier:
  202.                 cv2.line(vis, (x1+(w1+w2), y1), (x2+w1-w3, y2), acousticLines)
  203.  
  204.     cv2.imshow(win, vis)
  205.  
  206.  
  207. def draw_matches(window_name, kp_pairs1, kp_pairs2, electricImg, acousticImg,
  208.         mainImg, electricColorImg, acousticColorImg, mainColorImg):
  209.    
  210.     mkp1_1, mkp2_1 = zip(*kp_pairs1)
  211.     p1_1 = np.float32([kp.pt for kp in mkp1_1])
  212.     p2_1 = np.float32([kp.pt for kp in mkp2_1])
  213.  
  214.     mkp1_2, mkp2_2 = zip(*kp_pairs2)
  215.     p1_2 = np.float32([kp.pt for kp in mkp1_2])
  216.     p2_2 = np.float32([kp.pt for kp in mkp2_2])
  217.  
  218.     if len(kp_pairs1) >= 5:
  219.         H1, status1 = cv2.findHomography(p1_1, p2_1, cv2.RANSAC, 5.0)
  220.     else:
  221.         H1, status1 = None, None
  222.  
  223.     if len(kp_pairs2) >= 5:
  224.         H2, status2 = cv2.findHomography(p1_2, p2_2, cv2.RANSAC, 5.0)
  225.     else:
  226.         H2, status2 = None, None
  227.    
  228.     if len(p1_1):
  229.         explore_match(window_name, kp_pairs1, kp_pairs2, electricImg,
  230.         acousticImg, mainImg, electricColorImg, acousticColorImg,
  231.         mainColorImg, status1, H1, status2, H2)
  232.  
  233.  
  234.  
  235. def captureCamera(fn1, fn2, method):
  236.     cap = cv2.VideoCapture(0)
  237.     electricImg = cv2.imread(fn1, 0)
  238.     acousticImg = cv2.imread(fn2, 0)
  239.     electricColorImg = cv2.imread(fn1)
  240.     acousticColorImg = cv2.imread(fn2)
  241.     if electricImg is None:
  242.         print ('Failed to load fn1:', fn1)
  243.         sys.exit(1)
  244.    
  245.     if acousticImg is None:
  246.         print ('Failed to load fn2:', fn2)
  247.         sys.exit(1)
  248.  
  249.     while(True):
  250.         # Capture frame-by-frame
  251.         ret, frame = cap.read()
  252.  
  253.         mainImg = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)      
  254.         mainColorImg = frame
  255.  
  256.         kp_pairs1 = match_images(electricImg, mainImg, method)
  257.         kp_pairs2 = match_images(acousticImg, mainImg, method)
  258.    
  259.         if kp_pairs1:
  260.             draw_matches('find_obj', kp_pairs1, kp_pairs2, electricImg, acousticImg,
  261.             mainImg, electricColorImg, acousticColorImg, mainColorImg)  
  262.         else:
  263.             cv2.imshow('frame',frame)
  264.  
  265.         if cv2.waitKey(1) & 0xFF == ord('q'):
  266.             break
  267.  
  268.     cap.release()
  269.     cv2.destroyAllWindows()
  270.  
  271. #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  272. # Main
  273.  
  274.  
  275. def main():
  276.     if len(sys.argv) < 3:
  277.         print ("No filenames specified")
  278.         print ("USAGE: find_obj.py <image1> <image2>")
  279.         sys.exit(1)
  280.  
  281.     if (sys.argv[2] == "true"):
  282.         fn1 = "electricNarrow.png"
  283.         fn2 = "acousticNarrow.png"
  284.     elif (sys.argv[2] == "false"):
  285.         fn1 = "electricNoNeck.png"
  286.         fn2 = "acousticNoNeck.png"
  287.     else:
  288.         print ("Wrong parameter")
  289.  
  290.     fn3 = sys.argv[1]
  291.     method = sys.argv[3]
  292.  
  293.     if (fn3 == 'camera'):
  294.         captureCamera(fn1, fn2, method)
  295.     else:
  296.         findGuitars(fn1, fn2, fn3, method)
  297.         cv2.waitKey()
  298.         cv2.destroyAllWindows()
  299.  
  300. if __name__ == '__main__': main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement