Spaces:
Sleeping
Sleeping
File size: 2,962 Bytes
b69e47e 9bd3704 b69e47e 9bd3704 b69e47e 9bd3704 b69e47e 9bd3704 b69e47e 9bd3704 b69e47e 9bd3704 b69e47e 9bd3704 b69e47e 9bd3704 |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import os
import subprocess
import logging
from ultralytics import YOLO
from config import CONFIG
# Setup logging
logging.basicConfig(
level=logging.INFO,
format='{"time": "%(asctime)s", "level": "%(levelname)s", "message": "%(message)s"}'
)
logger = logging.getLogger(__name__)
def check_ffmpeg():
"""Verify FFmpeg is installed."""
try:
subprocess.run(["ffmpeg", "-version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
logger.info("FFmpeg is available.")
return True
except (subprocess.CalledProcessError, FileNotFoundError):
logger.error("FFmpeg is not installed or not found in PATH.")
return False
def check_model_files():
"""Verify model files exist or can be downloaded."""
model_path = CONFIG["MODEL_PATH"]
fallback_model = CONFIG["FALLBACK_MODEL"]
if not os.path.isfile(model_path):
logger.warning(f"Custom model {model_path} not found. Checking fallback.")
if not os.path.isfile(fallback_model):
logger.info(f"Downloading fallback model: {fallback_model}")
try:
import torch
torch.hub.download_url_to_file(
'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt',
fallback_model
)
logger.info(f"Downloaded {fallback_model}")
except Exception as e:
logger.error(f"Failed to download {fallback_model}: {e}")
return False
else:
logger.info(f"Custom model found: {model_path}")
return True
def test_model_inference():
"""Test YOLO model inference with a dummy image."""
try:
model_path = CONFIG["MODEL_PATH"] if os.path.isfile(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
model = YOLO(model_path)
# Create a dummy image (1x1 pixel)
import numpy as np
dummy_image = np.zeros((1, 1, 3), dtype=np.uint8)
results = model(dummy_image, conf=0.1, verbose=False)
logger.info("Model inference test passed.")
return True
except Exception as e:
logger.error(f"Model inference test failed: {e}")
return False
def run_integration_tests():
"""Run all integration tests."""
logger.info("Starting integration tests...")
tests = [
("FFmpeg availability", check_ffmpeg),
("Model files check", check_model_files),
("Model inference test", test_model_inference)
]
all_passed = True
for test_name, test_func in tests:
logger.info(f"Running test: {test_name}")
result = test_func()
if not result:
logger.error(f"Test failed: {test_name}")
all_passed = False
else:
logger.info(f"Test passed: {test_name}")
return all_passed
if __name__ == "__main__":
success = run_integration_tests()
exit(0 if success else 1) |