File size: 8,057 Bytes
53290d4 d83a6e5 53290d4 d83a6e5 014be24 d83a6e5 53290d4 deb9ec9 53290d4 deb9ec9 53290d4 | 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 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 | #!/usr/bin/env python3
"""
FastAPI Backend for Cat vs Dog Classification
Provides endpoints for image upload and prediction
"""
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
import numpy as np
import joblib
import json
import os
import cv2
from PIL import Image
import io
from huggingface_hub import hf_hub_download
app = FastAPI(title="Cat vs Dog Classification API", version="1.0.0")
# Enable CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Mount static files
os.makedirs("static", exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")
# Global variables for models and artifacts
model = None
scaler = None
label_encoder = None
metadata = None
def load_models():
"""Load trained models and artifacts from Hugging Face Hub"""
global model, scaler, label_encoder, metadata
try:
# Model repository configuration
repo_id = os.getenv("HF_MODEL_REPO", "enigmaceo/svm-classification-cat-and-dog")
# Download and load best model (compressed)
model_path = hf_hub_download(repo_id, "svm_best_model.pkl.gz")
import gzip
import pickle
with gzip.open(model_path, 'rb') as f:
model = pickle.load(f)
# Download and load scaler
scaler_path = hf_hub_download(repo_id, "scaler.pkl")
scaler = joblib.load(scaler_path)
# Download and load label encoder
encoder_path = hf_hub_download(repo_id, "label_encoder.pkl")
label_encoder = joblib.load(encoder_path)
# Download and load metadata
metadata_path = hf_hub_download(repo_id, "metadata.json")
with open(metadata_path, 'r') as f:
metadata = json.load(f)
print(f"Model and artifacts loaded successfully from {repo_id}")
return True
except Exception as e:
print(f"Error loading models from Hugging Face: {e}")
return False
def extract_hog_features(image, pixels_per_cell=(8, 8)):
"""Extract HOG features from image"""
from skimage.feature import hog
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
features, hog_img = hog(
gray,
orientations=9,
pixels_per_cell=pixels_per_cell,
cells_per_block=(2, 2),
block_norm='L2-Hys',
visualize=True,
transform_sqrt=True
)
return features.astype(np.float32)
def extract_color_histogram(image, bins=32):
"""Extract color histogram features from HSV image"""
hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
hist_h = np.histogram(hsv[:,:,0], bins=bins, range=(0, 180))[0]
hist_s = np.histogram(hsv[:,:,1], bins=bins, range=(0, 256))[0]
hist_v = np.histogram(hsv[:,:,2], bins=bins, range=(0, 256))[0]
return np.concatenate([hist_h, hist_s, hist_v]).astype(np.float32)
def extract_lbp_features(image, radius=3, n_points=24):
"""Extract Local Binary Pattern features for texture"""
from skimage.feature import local_binary_pattern
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
lbp = local_binary_pattern(gray, n_points, radius, method='uniform')
hist, _ = np.histogram(lbp.ravel(), bins=n_points + 2)
hist = hist.astype(np.float32)
hist /= (hist.sum() + 1e-7) # Normalize
return hist
def extract_features_from_image(image_data: bytes) -> np.ndarray:
"""
Extract HOG, color histogram, and LBP features from uploaded image
Same feature extraction as used in training
"""
try:
# Convert bytes to PIL Image
image = Image.open(io.BytesIO(image_data))
# Convert to numpy array and RGB
img_array = np.array(image)
if len(img_array.shape) == 2: # Grayscale
img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
elif img_array.shape[2] == 4: # RGBA
img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
# Resize to match training size
img_resized = cv2.resize(img_array, (128, 128))
# Extract HOG features
hog_feat = extract_hog_features(img_resized)
# Extract color histogram
col_feat = extract_color_histogram(img_resized)
# Extract LBP features
lbp_feat = extract_lbp_features(img_resized)
# Combine features
combined_features = np.concatenate([hog_feat, col_feat, lbp_feat])
return combined_features.reshape(1, -1)
except Exception as e:
raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")
@app.on_event("startup")
async def startup_event():
"""Load models on startup"""
success = load_models()
if not success:
print("Warning: Could not load models. Please run training script first.")
@app.get("/", response_class=HTMLResponse)
async def root():
"""Serve the dashboard"""
try:
with open("dashboard/index.html", "r") as f:
return HTMLResponse(content=f.read())
except FileNotFoundError:
return HTMLResponse(content="<h1>Cat vs Dog Classification</h1><p>Dashboard not found. Please check dashboard folder.</p>")
@app.get("/api/health")
async def health_check():
"""Health check endpoint"""
return {
"status": "healthy",
"model_loaded": model is not None,
"best_kernel": metadata.get('best_kernel') if metadata else None
}
@app.post("/api/predict")
async def predict_image(file: UploadFile = File(...)):
"""Predict cat or dog from uploaded image"""
if not model:
raise HTTPException(status_code=503, detail="Model not loaded")
if not scaler:
raise HTTPException(status_code=503, detail="Scaler not loaded")
if not label_encoder:
raise HTTPException(status_code=503, detail="Label encoder not loaded")
try:
# Read image data
image_data = await file.read()
# Extract features
features = extract_features_from_image(image_data)
features_scaled = scaler.transform(features)
# Extract individual features for display
img_array = np.array(Image.open(io.BytesIO(image_data)))
if len(img_array.shape) == 2: # Grayscale
img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
elif img_array.shape[2] == 4: # RGBA
img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
img_resized = cv2.resize(img_array, (128, 128))
hog_feat = extract_hog_features(img_resized)
col_feat = extract_color_histogram(img_resized)
lbp_feat = extract_lbp_features(img_resized)
# Make prediction
prediction = model.predict(features_scaled)[0]
class_id = int(prediction)
class_name = label_encoder.inverse_transform([class_id])[0]
# Get confidence if available
confidence = None
if hasattr(model, 'decision_function'):
decision_values = model.decision_function(features_scaled)[0]
exp_values = np.exp(decision_values - np.max(decision_values))
probabilities = exp_values / np.sum(exp_values)
confidence = float(np.max(probabilities))
return {
"prediction": {
"class_id": class_id,
"class_name": class_name,
"confidence": confidence
},
"features": {
"hog_size": len(hog_feat),
"color_size": len(col_feat),
"lbp_size": len(lbp_feat)
}
}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
|