Register
Login
Resources
Docs Blog Datasets Glossary Case Studies Tutorials & Webinars
Product
Data Engine LLMs Platform Enterprise
Pricing Explore
Connect to our Discord channel

test_explorer.py 2.5 KB

You have to be logged in to leave a comment. Sign In
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. import PIL
  3. import pytest
  4. from ultralytics import Explorer
  5. from ultralytics.utils import ASSETS
  6. from ultralytics.utils.torch_utils import TORCH_1_13
  7. @pytest.mark.slow
  8. @pytest.mark.skipif(not TORCH_1_13, reason="Explorer requires torch>=1.13")
  9. def test_similarity():
  10. """Test the correctness and response length of similarity calculations and SQL queries in the Explorer."""
  11. exp = Explorer(data="coco8.yaml")
  12. exp.create_embeddings_table()
  13. similar = exp.get_similar(idx=1)
  14. assert len(similar) == 4
  15. similar = exp.get_similar(img=ASSETS / "bus.jpg")
  16. assert len(similar) == 4
  17. similar = exp.get_similar(idx=[1, 2], limit=2)
  18. assert len(similar) == 2
  19. sim_idx = exp.similarity_index()
  20. assert len(sim_idx) == 4
  21. sql = exp.sql_query("WHERE labels LIKE '%zebra%'")
  22. assert len(sql) == 1
  23. @pytest.mark.slow
  24. @pytest.mark.skipif(not TORCH_1_13, reason="Explorer requires torch>=1.13")
  25. def test_det():
  26. """Test detection functionalities and verify embedding table includes bounding boxes."""
  27. exp = Explorer(data="coco8.yaml", model="yolov8n.pt")
  28. exp.create_embeddings_table(force=True)
  29. assert len(exp.table.head()["bboxes"]) > 0
  30. similar = exp.get_similar(idx=[1, 2], limit=10)
  31. assert len(similar) > 0
  32. # This is a loose test, just checks errors not correctness
  33. similar = exp.plot_similar(idx=[1, 2], limit=10)
  34. assert isinstance(similar, PIL.Image.Image)
  35. @pytest.mark.slow
  36. @pytest.mark.skipif(not TORCH_1_13, reason="Explorer requires torch>=1.13")
  37. def test_seg():
  38. """Test segmentation functionalities and ensure the embedding table includes segmentation masks."""
  39. exp = Explorer(data="coco8-seg.yaml", model="yolov8n-seg.pt")
  40. exp.create_embeddings_table(force=True)
  41. assert len(exp.table.head()["masks"]) > 0
  42. similar = exp.get_similar(idx=[1, 2], limit=10)
  43. assert len(similar) > 0
  44. similar = exp.plot_similar(idx=[1, 2], limit=10)
  45. assert isinstance(similar, PIL.Image.Image)
  46. @pytest.mark.slow
  47. @pytest.mark.skipif(not TORCH_1_13, reason="Explorer requires torch>=1.13")
  48. def test_pose():
  49. """Test pose estimation functionality and verify the embedding table includes keypoints."""
  50. exp = Explorer(data="coco8-pose.yaml", model="yolov8n-pose.pt")
  51. exp.create_embeddings_table(force=True)
  52. assert len(exp.table.head()["keypoints"]) > 0
  53. similar = exp.get_similar(idx=[1, 2], limit=10)
  54. assert len(similar) > 0
  55. similar = exp.plot_similar(idx=[1, 2], limit=10)
  56. assert isinstance(similar, PIL.Image.Image)
Tip!

Press p or to see the previous file or, n or to see the next file

Comments

Loading...