helper.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155
  1. import cv2
  2. import numpy as np
  3. from io import BytesIO
  4. from PIL import Image
  5. from ppadb.client import Client as AdbClient
  6. def get_current_screen():
  7. return current_screen
  8. def capture_current_screen():
  9. global current_screen
  10. current_screen = device.screencap()
  11. return current_screen
  12. def find_center(x1, y1, x2, y2):
  13. centerX = round(x1 + (x2 - x1) / 2)
  14. centerY = round(y1 + (y2 - y1) / 2)
  15. return centerX, centerY
  16. def tap(x, y=None):
  17. # Check if x is an int
  18. if isinstance(x, int):
  19. if not isinstance(y, int):
  20. raise ValueError("y must be an int when x is an int")
  21. # Construct the location string from both x and y
  22. location = f"{x} {y}"
  23. # Check if x is a string
  24. elif isinstance(x, str):
  25. location = x
  26. else:
  27. raise TypeError("x must be either an int or a string")
  28. # Assuming 'device' is a previously defined object with a 'shell' method
  29. action = f"input tap {location}"
  30. print(action)
  31. device.shell(action)
  32. def tap_button(template):
  33. button = find_template(template)
  34. if len(button) == 0:
  35. return
  36. tap(f"{button[0][0]} {button[0][1]}")
  37. def swipe(start, end, duration=1000):
  38. action = f"input swipe {start} {end} {duration}"
  39. print(action)
  40. device.shell(action)
  41. def look_for_templates(templates):
  42. for name, template in templates.items():
  43. locations = find_template(template)
  44. if len(locations) > 0:
  45. return name, locations
  46. return None, None
  47. def find_template(template_image):
  48. target_image = Image.open(BytesIO(get_current_screen()))
  49. # Convert the image to a NumPy array and then to BGR format (which OpenCV uses)
  50. target_image = np.array(target_image)
  51. target_image = cv2.cvtColor(target_image, cv2.COLOR_RGB2BGR)
  52. h, w = template_image.shape[:-1]
  53. # Template matching
  54. result = cv2.matchTemplate(target_image, template_image, cv2.TM_CCOEFF_NORMED)
  55. # Define a threshold
  56. threshold = 0.9 # Adjust this threshold based on your requirements
  57. # Finding all locations where match exceeds threshold
  58. locations = np.where(result >= threshold)
  59. locations = list(zip(*locations[::-1]))
  60. # Create list of rectangles
  61. rectangles = [(*loc, loc[0] + w, loc[1] + h) for loc in locations]
  62. # Apply non-maximum suppression to remove overlaps
  63. rectangles = non_max_suppression(rectangles, 0.3)
  64. # Initialize an empty list to store coordinates
  65. coordinates = []
  66. for startX, startY, endX, endY in rectangles:
  67. # Append the coordinate pair to the list
  68. coordinates.append(find_center(startX, startY, endX, endY))
  69. # Sort the coordinates by y value in ascending order
  70. return sorted(coordinates, key=lambda x: x[1])
  71. def non_max_suppression(boxes, overlapThresh):
  72. if len(boxes) == 0:
  73. return []
  74. # Convert to float
  75. boxes = np.array(boxes, dtype="float")
  76. # Initialize the list of picked indexes
  77. pick = []
  78. # Grab the coordinates of the bounding boxes
  79. x1 = boxes[:, 0]
  80. y1 = boxes[:, 1]
  81. x2 = boxes[:, 2]
  82. y2 = boxes[:, 3]
  83. # Compute the area of the bounding boxes and sort by bottom-right y-coordinate
  84. area = (x2 - x1 + 1) * (y2 - y1 + 1)
  85. idxs = np.argsort(y2)
  86. # Keep looping while some indexes still remain in the indexes list
  87. while len(idxs) > 0:
  88. # Grab the last index in the indexes list and add the index value to the list of picked indexes
  89. last = len(idxs) - 1
  90. i = idxs[last]
  91. pick.append(i)
  92. # Find the largest (x, y) coordinates for the start of the bounding box and the smallest (x, y)
  93. # coordinates for the end of the bounding box
  94. xx1 = np.maximum(x1[i], x1[idxs[:last]])
  95. yy1 = np.maximum(y1[i], y1[idxs[:last]])
  96. xx2 = np.minimum(x2[i], x2[idxs[:last]])
  97. yy2 = np.minimum(y2[i], y2[idxs[:last]])
  98. # Compute the width and height of the bounding box
  99. w = np.maximum(0, xx2 - xx1 + 1)
  100. h = np.maximum(0, yy2 - yy1 + 1)
  101. # Compute the ratio of overlap
  102. overlap = (w * h) / area[idxs[:last]]
  103. # Delete all indexes from the index list that have overlap greater than the threshold
  104. idxs = np.delete(
  105. idxs, np.concatenate(([last], np.where(overlap > overlapThresh)[0]))
  106. )
  107. # Return only the bounding boxes that were picked
  108. return boxes[pick].astype("int")
  109. client = AdbClient(host="127.0.0.1", port=5037)
  110. device = client.device("192.168.178.32:5555")
  111. current_screen = capture_current_screen()