Spaces:
Sleeping
Sleeping
| 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) |