locate.py 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. import cv2
  2. import numpy as np
  3. def non_max_suppression(boxes, overlapThresh):
  4. if len(boxes) == 0:
  5. return []
  6. # Convert to float
  7. boxes = np.array(boxes, dtype="float")
  8. # Initialize the list of picked indexes
  9. pick = []
  10. # Grab the coordinates of the bounding boxes
  11. x1 = boxes[:,0]
  12. y1 = boxes[:,1]
  13. x2 = boxes[:,2]
  14. y2 = boxes[:,3]
  15. # Compute the area of the bounding boxes and sort by bottom-right y-coordinate
  16. area = (x2 - x1 + 1) * (y2 - y1 + 1)
  17. idxs = np.argsort(y2)
  18. # Keep looping while some indexes still remain in the indexes list
  19. while len(idxs) > 0:
  20. # Grab the last index in the indexes list and add the index value to the list of picked indexes
  21. last = len(idxs) - 1
  22. i = idxs[last]
  23. pick.append(i)
  24. # Find the largest (x, y) coordinates for the start of the bounding box and the smallest (x, y)
  25. # coordinates for the end of the bounding box
  26. xx1 = np.maximum(x1[i], x1[idxs[:last]])
  27. yy1 = np.maximum(y1[i], y1[idxs[:last]])
  28. xx2 = np.minimum(x2[i], x2[idxs[:last]])
  29. yy2 = np.minimum(y2[i], y2[idxs[:last]])
  30. # Compute the width and height of the bounding box
  31. w = np.maximum(0, xx2 - xx1 + 1)
  32. h = np.maximum(0, yy2 - yy1 + 1)
  33. # Compute the ratio of overlap
  34. overlap = (w * h) / area[idxs[:last]]
  35. # Delete all indexes from the index list that have overlap greater than the threshold
  36. idxs = np.delete(idxs, np.concatenate(([last], np.where(overlap > overlapThresh)[0])))
  37. # Return only the bounding boxes that were picked
  38. return boxes[pick].astype("int")
  39. # Load images
  40. target_image = cv2.imread('test/20240108_021855.png')
  41. template_image = cv2.imread('templates/end_of_log.jpg')
  42. w, h = template_image.shape[:-1]
  43. # Template matching
  44. result = cv2.matchTemplate(target_image, template_image, cv2.TM_CCOEFF_NORMED)
  45. # Define a threshold
  46. threshold = 0.9 # Adjust this threshold based on your requirements
  47. # Finding all locations where match exceeds threshold
  48. locations = np.where(result >= threshold)
  49. locations = list(zip(*locations[::-1]))
  50. # Create list of rectangles
  51. rectangles = [(*loc, loc[0] + w, loc[1] + h) for loc in locations]
  52. # Apply non-maximum suppression to remove overlaps
  53. rectangles = non_max_suppression(rectangles, 0.3)
  54. # Draw rectangles around matches
  55. for (startX, startY, endX, endY) in rectangles:
  56. cv2.rectangle(target_image, (startX, startY), (endX, endY), (0, 255, 0), 2)
  57. # Print the coordinates of the rectangle
  58. centerX = round(startX + (endX-startX)/2)
  59. centerY = round(startY + (endY-startY)/2)
  60. print(f"center coordinates: {centerX}/{centerY}")
  61. # Display the result
  62. cv2.imshow('Matched Results', target_image)
  63. cv2.waitKey(0)
  64. cv2.destroyAllWindows()