Spaces:
Sleeping
Sleeping
| import torch | |
| from src.model import CatDogClassifier | |
| from src.config import CatDogClassifierConfigs | |
| def inference_pipeline( | |
| image_path: str = "datasets/single_prediction/cat_or_dog_1.jpg", | |
| model_path: str = "checkpoints/ckpt_23_10_2025/best_cat_dog_classifier_model_20251019_122336.pth" | |
| ): | |
| # Initialize model | |
| model_configs = CatDogClassifierConfigs( | |
| device="cuda" if torch.cuda.is_available() else "cpu", | |
| input_channels=3, | |
| num_classes=2, | |
| learning_rate=0.001, | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| num_layers=3, | |
| use_amp=False | |
| ) | |
| # Load state_dict | |
| model = CatDogClassifier(configs=model_configs) | |
| # Load state_dict (both local & remote) | |
| state_dict = torch.load(model_path, map_location=model_configs.device) | |
| model.load_state_dict(state_dict) | |
| model.eval() | |
| y_pred = model.predict( | |
| model=model, | |
| image_path=image_path | |
| ) | |
| print(f"Predicted class for the image {image_path}: {y_pred}") | |
| return y_pred | |
| if __name__ == "__main__": | |
| y_pred = inference_pipeline("D:\\Desktop\\stores\\Application\\GoldenOwl\\technical_test\\test_image_2.jpg") | |
| print(y_pred) |