Kye
1 year ago
parent
3c0a12b4ae
commit
acdcdd61a6
@ -1,394 +0,0 @@
|
||||
import pytest
|
||||
import os
|
||||
from PIL import Image
|
||||
from swarms.models.kosmos2 import Kosmos2, Detections
|
||||
|
||||
|
||||
# Fixture for a sample image
|
||||
@pytest.fixture
|
||||
def sample_image():
|
||||
image = Image.new("RGB", (224, 224))
|
||||
return image
|
||||
|
||||
|
||||
# Fixture for initializing Kosmos2
|
||||
@pytest.fixture
|
||||
def kosmos2():
|
||||
return Kosmos2.initialize()
|
||||
|
||||
|
||||
# Test Kosmos2 initialization
|
||||
def test_kosmos2_initialization(kosmos2):
|
||||
assert kosmos2 is not None
|
||||
|
||||
|
||||
# Test Kosmos2 with a sample image
|
||||
def test_kosmos2_with_sample_image(kosmos2, sample_image):
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
# Mocked extract_entities function for testing
|
||||
def mock_extract_entities(text):
|
||||
return [
|
||||
("entity1", (0.1, 0.2, 0.3, 0.4)),
|
||||
("entity2", (0.5, 0.6, 0.7, 0.8)),
|
||||
]
|
||||
|
||||
|
||||
# Mocked process_entities_to_detections function for testing
|
||||
def mock_process_entities_to_detections(entities, image):
|
||||
return Detections(
|
||||
xyxy=[(10, 20, 30, 40), (50, 60, 70, 80)],
|
||||
class_id=[0, 0],
|
||||
confidence=[1.0, 1.0],
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with mocked entity extraction and detection
|
||||
def test_kosmos2_with_mocked_extraction_and_detection(
|
||||
kosmos2, sample_image, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
kosmos2, "extract_entities", mock_extract_entities
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
kosmos2,
|
||||
"process_entities_to_detections",
|
||||
mock_process_entities_to_detections,
|
||||
)
|
||||
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 2
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with empty entity extraction
|
||||
def test_kosmos2_with_empty_extraction(
|
||||
kosmos2, sample_image, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr(kosmos2, "extract_entities", lambda x: [])
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with invalid image path
|
||||
def test_kosmos2_with_invalid_image_path(kosmos2):
|
||||
with pytest.raises(Exception):
|
||||
kosmos2(img="invalid_image_path.jpg")
|
||||
|
||||
|
||||
# Additional tests can be added for various scenarios and edge cases
|
||||
|
||||
|
||||
# Test Kosmos2 with a larger image
|
||||
def test_kosmos2_with_large_image(kosmos2):
|
||||
large_image = Image.new("RGB", (1024, 768))
|
||||
detections = kosmos2(img=large_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with different image formats
|
||||
def test_kosmos2_with_different_image_formats(kosmos2, tmp_path):
|
||||
# Create a temporary directory
|
||||
temp_dir = tmp_path / "images"
|
||||
temp_dir.mkdir()
|
||||
|
||||
# Create sample images in different formats
|
||||
image_formats = ["jpeg", "png", "gif", "bmp"]
|
||||
for format in image_formats:
|
||||
image_path = temp_dir / f"sample_image.{format}"
|
||||
Image.new("RGB", (224, 224)).save(image_path)
|
||||
|
||||
# Test Kosmos2 with each image format
|
||||
for format in image_formats:
|
||||
image_path = temp_dir / f"sample_image.{format}"
|
||||
detections = kosmos2(img=image_path)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with a non-existent model
|
||||
def test_kosmos2_with_non_existent_model(kosmos2):
|
||||
with pytest.raises(Exception):
|
||||
kosmos2.model = None
|
||||
kosmos2(img="sample_image.jpg")
|
||||
|
||||
|
||||
# Test Kosmos2 with a non-existent processor
|
||||
def test_kosmos2_with_non_existent_processor(kosmos2):
|
||||
with pytest.raises(Exception):
|
||||
kosmos2.processor = None
|
||||
kosmos2(img="sample_image.jpg")
|
||||
|
||||
|
||||
# Test Kosmos2 with missing image
|
||||
def test_kosmos2_with_missing_image(kosmos2):
|
||||
with pytest.raises(Exception):
|
||||
kosmos2(img="non_existent_image.jpg")
|
||||
|
||||
|
||||
# ... (previous tests)
|
||||
|
||||
|
||||
# Test Kosmos2 with a non-existent model and processor
|
||||
def test_kosmos2_with_non_existent_model_and_processor(kosmos2):
|
||||
with pytest.raises(Exception):
|
||||
kosmos2.model = None
|
||||
kosmos2.processor = None
|
||||
kosmos2(img="sample_image.jpg")
|
||||
|
||||
|
||||
# Test Kosmos2 with a corrupted image
|
||||
def test_kosmos2_with_corrupted_image(kosmos2, tmp_path):
|
||||
# Create a temporary directory
|
||||
temp_dir = tmp_path / "images"
|
||||
temp_dir.mkdir()
|
||||
|
||||
# Create a corrupted image
|
||||
corrupted_image_path = temp_dir / "corrupted_image.jpg"
|
||||
with open(corrupted_image_path, "wb") as f:
|
||||
f.write(b"corrupted data")
|
||||
|
||||
with pytest.raises(Exception):
|
||||
kosmos2(img=corrupted_image_path)
|
||||
|
||||
|
||||
# Test Kosmos2 with a large batch size
|
||||
def test_kosmos2_with_large_batch_size(kosmos2, sample_image):
|
||||
kosmos2.batch_size = 32
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with an invalid compute type
|
||||
def test_kosmos2_with_invalid_compute_type(kosmos2, sample_image):
|
||||
kosmos2.compute_type = "invalid_compute_type"
|
||||
with pytest.raises(Exception):
|
||||
kosmos2(img=sample_image)
|
||||
|
||||
|
||||
# Test Kosmos2 with a valid HF API key
|
||||
def test_kosmos2_with_valid_hf_api_key(kosmos2, sample_image):
|
||||
kosmos2.hf_api_key = "valid_api_key"
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 2
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with an invalid HF API key
|
||||
def test_kosmos2_with_invalid_hf_api_key(kosmos2, sample_image):
|
||||
kosmos2.hf_api_key = "invalid_api_key"
|
||||
with pytest.raises(Exception):
|
||||
kosmos2(img=sample_image)
|
||||
|
||||
|
||||
# Test Kosmos2 with a very long generated text
|
||||
def test_kosmos2_with_long_generated_text(
|
||||
kosmos2, sample_image, monkeypatch
|
||||
):
|
||||
def mock_generate_text(*args, **kwargs):
|
||||
return "A" * 10000
|
||||
|
||||
monkeypatch.setattr(kosmos2.model, "generate", mock_generate_text)
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with entities containing special characters
|
||||
def test_kosmos2_with_entities_containing_special_characters(
|
||||
kosmos2, sample_image, monkeypatch
|
||||
):
|
||||
def mock_extract_entities(text):
|
||||
return [
|
||||
(
|
||||
"entity1 with special characters (ü, ö, etc.)",
|
||||
(0.1, 0.2, 0.3, 0.4),
|
||||
)
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
kosmos2, "extract_entities", mock_extract_entities
|
||||
)
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 1
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with image containing multiple objects
|
||||
def test_kosmos2_with_image_containing_multiple_objects(
|
||||
kosmos2, sample_image, monkeypatch
|
||||
):
|
||||
def mock_extract_entities(text):
|
||||
return [
|
||||
("entity1", (0.1, 0.2, 0.3, 0.4)),
|
||||
("entity2", (0.5, 0.6, 0.7, 0.8)),
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
kosmos2, "extract_entities", mock_extract_entities
|
||||
)
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 2
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with image containing no objects
|
||||
def test_kosmos2_with_image_containing_no_objects(
|
||||
kosmos2, sample_image, monkeypatch
|
||||
):
|
||||
def mock_extract_entities(text):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(
|
||||
kosmos2, "extract_entities", mock_extract_entities
|
||||
)
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with a valid YouTube video URL
|
||||
def test_kosmos2_with_valid_youtube_video_url(kosmos2):
|
||||
youtube_video_url = "https://www.youtube.com/watch?v=VIDEO_ID"
|
||||
detections = kosmos2(video_url=youtube_video_url)
|
||||
assert isinstance(detections, Detections)
|
||||
assert (
|
||||
len(detections.xyxy)
|
||||
== len(detections.class_id)
|
||||
== len(detections.confidence)
|
||||
== 2
|
||||
)
|
||||
|
||||
|
||||
# Test Kosmos2 with an invalid YouTube video URL
|
||||
def test_kosmos2_with_invalid_youtube_video_url(kosmos2):
|
||||
invalid_youtube_video_url = (
|
||||
"https://www.youtube.com/invalid_video"
|
||||
)
|
||||
with pytest.raises(Exception):
|
||||
kosmos2(video_url=invalid_youtube_video_url)
|
||||
|
||||
|
||||
# Test Kosmos2 with no YouTube video URL provided
|
||||
def test_kosmos2_with_no_youtube_video_url(kosmos2):
|
||||
with pytest.raises(Exception):
|
||||
kosmos2(video_url=None)
|
||||
|
||||
|
||||
# Test Kosmos2 installation
|
||||
def test_kosmos2_installation():
|
||||
kosmos2 = Kosmos2()
|
||||
kosmos2.install()
|
||||
assert os.path.exists("video.mp4")
|
||||
assert os.path.exists("video.mp3")
|
||||
os.remove("video.mp4")
|
||||
os.remove("video.mp3")
|
||||
|
||||
|
||||
# Test Kosmos2 termination
|
||||
def test_kosmos2_termination(kosmos2):
|
||||
kosmos2.terminate()
|
||||
assert kosmos2.process is None
|
||||
|
||||
|
||||
# Test Kosmos2 start_process method
|
||||
def test_kosmos2_start_process(kosmos2):
|
||||
kosmos2.start_process()
|
||||
assert kosmos2.process is not None
|
||||
|
||||
|
||||
# Test Kosmos2 preprocess_code method
|
||||
def test_kosmos2_preprocess_code(kosmos2):
|
||||
code = "print('Hello, World!')"
|
||||
preprocessed_code = kosmos2.preprocess_code(code)
|
||||
assert isinstance(preprocessed_code, str)
|
||||
assert "end_of_execution" in preprocessed_code
|
||||
|
||||
|
||||
# Test Kosmos2 run method with debug mode
|
||||
def test_kosmos2_run_with_debug_mode(kosmos2, sample_image):
|
||||
kosmos2.debug_mode = True
|
||||
detections = kosmos2(img=sample_image)
|
||||
assert isinstance(detections, Detections)
|
||||
|
||||
|
||||
# Test Kosmos2 handle_stream_output method
|
||||
def test_kosmos2_handle_stream_output(kosmos2):
|
||||
stream_output = "Sample output"
|
||||
kosmos2.handle_stream_output(stream_output, is_error=False)
|
||||
|
||||
|
||||
# Test Kosmos2 run method with invalid image path
|
||||
def test_kosmos2_run_with_invalid_image_path(kosmos2):
|
||||
with pytest.raises(Exception):
|
||||
kosmos2.run(img="invalid_image_path.jpg")
|
||||
|
||||
|
||||
# Test Kosmos2 run method with invalid video URL
|
||||
def test_kosmos2_run_with_invalid_video_url(kosmos2):
|
||||
with pytest.raises(Exception):
|
||||
kosmos2.run(video_url="invalid_video_url")
|
||||
|
||||
|
||||
# ... (more tests)
|
||||
Loading…
Reference in new issue