vikenkd's picture
[feat]: upload files
7fab19a
Raw
History Blame Contribute Delete
1.2 kB
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)