Spaces:
Sleeping
Sleeping
File size: 4,708 Bytes
d2a2955 |
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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
"""
Model inference module for QuickDraw sketch classification.
Handles model loading and prediction logic.
"""
import os
import numpy as np
import tensorflow as tf
from typing import List, Dict
import logging
logger = logging.getLogger(__name__)
class SketchClassifier:
"""QuickDraw sketch classifier"""
def __init__(self, model_path: str = None):
"""
Initialize the classifier with a trained model.
Args:
model_path: Path to the trained model file. If None, uses default path.
"""
# Extended class list matching Model-Training.py
self.class_names = [
# Animals
"cat", "dog", "bird", "fish", "bear", "butterfly", "bee", "spider",
# Buildings & Structures
"house", "castle", "barn", "bridge", "lighthouse", "church",
# Transportation
"car", "airplane", "bicycle", "boat", "train", "truck", "bus",
# Nature
"tree", "flower", "sun", "moon", "cloud", "mountain", "river",
# Common Objects
"apple", "banana", "book", "chair", "table", "cup", "umbrella",
# People & Body
"face", "eye", "hand", "foot",
# Shapes & Symbols
"circle", "triangle", "square", "star", "heart",
# Tools & Items
"sword", "axe", "hammer", "key", "crown"
]
# Default model path
if model_path is None:
model_path = os.path.join("saved_models", "quickdraw_house_cat_dog_car.keras")
# Check if model exists
if not os.path.exists(model_path):
# Try .h5 format as fallback
h5_path = model_path.replace(".keras", ".h5")
if os.path.exists(h5_path):
model_path = h5_path
logger.info(f"Using H5 model format: {model_path}")
else:
raise FileNotFoundError(
f"Model file not found at {model_path}. "
"Please train the model first using Model-Training.py"
)
logger.info(f"Loading model from: {model_path}")
self.model = tf.keras.models.load_model(model_path)
logger.info("Model loaded successfully!")
# Verify input shape
self.input_shape = self.model.input_shape[1:] # (28, 28, 1)
logger.info(f"Model input shape: {self.input_shape}")
def predict(self, image: np.ndarray, top_k: int = 3) -> List[Dict[str, any]]:
"""
Make prediction on a preprocessed image.
Args:
image: Preprocessed image array of shape (1, 28, 28, 1)
top_k: Number of top predictions to return
Returns:
List of dictionaries containing class names and confidence scores
"""
# Validate input shape
if image.shape != (1, 28, 28, 1):
raise ValueError(
f"Expected input shape (1, 28, 28, 1), got {image.shape}. "
"Please preprocess the image first."
)
# Make prediction
predictions = self.model.predict(image, verbose=0)
# Get top k predictions
top_indices = np.argsort(predictions[0])[::-1][:top_k]
results = []
for idx in top_indices:
results.append({
"class": self.class_names[idx],
"confidence": float(predictions[0][idx]),
"confidence_percent": f"{predictions[0][idx] * 100:.2f}%"
})
return results
def predict_batch(self, images: np.ndarray, top_k: int = 3) -> List[List[Dict[str, any]]]:
"""
Make predictions on a batch of preprocessed images.
Args:
images: Batch of preprocessed images of shape (N, 28, 28, 1)
top_k: Number of top predictions to return per image
Returns:
List of prediction results for each image
"""
# Make predictions
predictions = self.model.predict(images, verbose=0)
results = []
for pred in predictions:
# Get top k predictions for this image
top_indices = np.argsort(pred)[::-1][:top_k]
image_results = []
for idx in top_indices:
image_results.append({
"class": self.class_names[idx],
"confidence": float(pred[idx]),
"confidence_percent": f"{pred[idx] * 100:.2f}%"
})
results.append(image_results)
return results
|