Fix side case with Template Matching (#105)

* fix side cases with tm

* miinor change to readme

README.md

@@ -124,6 +124,7 @@ You can also add custom tools to the agent:
 
 ```python
 import vision_agent as va
+import numpy as np
 
 @va.tools.register_tool(imports=["import numpy as np"])
 def custom_tool(image_path: str) -> str:
@@ -140,7 +141,6 @@ def custom_tool(image_path: str) -> str:
  >>> custom_tool("image.jpg")
  """
 
- import numpy as np
  return np.zeros((10, 10))
 ```
 

examples/custom_tools/template_match.py

@@ -54,6 +54,11 @@ def template_matching_with_rotation(
  for angle in range(0, max_rotation, step):
  # Rotate the template
  rotated_template = rotate_image(template_gray, angle)
+ if (
+ rotated_template.shape[0] > main_image_gray.shape[0]
+ or rotated_template.shape[1] > main_image_gray.shape[1]
+ ):
+ continue
 
  # Perform template matching
  result = cv2.matchTemplate(
@@ -69,17 +74,18 @@ def template_matching_with_rotation(
  )
  scores.append(result[y, x])
 
- indices = (
- nms(
- torch.tensor(boxes).float(),
- torch.tensor(scores).float(),
- 0.2,
+ if len(boxes) > 0:
+ indices = (
+ nms(
+ torch.tensor(boxes).float(),
+ torch.tensor(scores).float(),
+ 0.2,
+ )
+ .numpy()
+ .tolist()
  )
- .numpy()
- .tolist()
- )
- boxes = [boxes[i] for i in indices]
- scores = [scores[i] for i in indices]
+ boxes = [boxes[i] for i in indices]
+ scores = [scores[i] for i in indices]
 
  if visualize:
  # Draw a rectangle around the best match