Spaces:
Running
Running
File size: 1,559 Bytes
713f590 | 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 | import os
import cv2
import numpy as np
import torch
from model import CatDogCNN
LABELS = {
0: "cat",
1: "dog"
}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CatDogCNN().to(device)
# Load model from the same directory as this file
_model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cat_dog_model.pth")
model.load_state_dict(torch.load(_model_path, map_location=device))
model.eval()
def preprocess_image(image):
if image is None:
raise ValueError("No image provided")
if hasattr(image, "convert"):
image = np.array(image.convert("RGB"))
else:
image = np.asarray(image)
if image.ndim == 2:
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
elif image.shape[-1] == 4:
image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
elif image.shape[-1] == 1:
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
image = cv2.resize(image, (64, 64))
image = image.astype(np.float32) / 255.0
image = np.transpose(image, (2, 0, 1))
image = torch.from_numpy(image).float().unsqueeze(0)
return image
def predict(image):
if image is None:
return {"cat": 0.0, "dog": 0.0}
image = preprocess_image(image)
image = image.to(device)
with torch.no_grad():
logits = model(image)
probabilities = torch.softmax(logits, dim=1)[0]
return {
"cat": float(probabilities[0]),
"dog": float(probabilities[1])
}
|