enigmaceo/svm
0
1#!/usr/bin/env python32"""3FastAPI Backend for Cat vs Dog Classification4Provides endpoints for image upload and prediction5"""6 7from fastapi import FastAPI, File, UploadFile, HTTPException8from fastapi.middleware.cors import CORSMiddleware9from fastapi.responses import HTMLResponse10from fastapi.staticfiles import StaticFiles11import numpy as np12import joblib13import json14import os15import cv216from PIL import Image17import io18from huggingface_hub import hf_hub_download19 20app = FastAPI(title="Cat vs Dog Classification API", version="1.0.0")21 22# Enable CORS23app.add_middleware(24 CORSMiddleware,25 allow_origins=["*"],26 allow_credentials=True,27 allow_methods=["*"],28 allow_headers=["*"],29)30 31# Mount static files32os.makedirs("static", exist_ok=True)33app.mount("/static", StaticFiles(directory="static"), name="static")34 35# Global variables for models and artifacts36model = None37scaler = None38label_encoder = None39metadata = None40 41def load_models():42 """Load trained models and artifacts from Hugging Face Hub"""43 global model, scaler, label_encoder, metadata44 45 try:46 # Model repository configuration47 repo_id = os.getenv("HF_MODEL_REPO", "enigmaceo/svm-classification-cat-and-dog")48 49 # Download and load best model (compressed)50 model_path = hf_hub_download(repo_id, "svm_best_model.pkl.gz")51 import gzip52 import pickle53 with gzip.open(model_path, 'rb') as f:54 model = pickle.load(f)55 56 # Download and load scaler57 scaler_path = hf_hub_download(repo_id, "scaler.pkl")58 scaler = joblib.load(scaler_path)59 60 # Download and load label encoder61 encoder_path = hf_hub_download(repo_id, "label_encoder.pkl")62 label_encoder = joblib.load(encoder_path)63 64 # Download and load metadata65 metadata_path = hf_hub_download(repo_id, "metadata.json")66 with open(metadata_path, 'r') as f:67 metadata = json.load(f)68 69 print(f"Model and artifacts loaded successfully from {repo_id}")70 return True71 except Exception as e:72 print(f"Error loading models from Hugging Face: {e}")73 return False74 75def extract_hog_features(image, pixels_per_cell=(8, 8)):76 """Extract HOG features from image"""77 from skimage.feature import hog78 gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)79 features, hog_img = hog(80 gray,81 orientations=9,82 pixels_per_cell=pixels_per_cell,83 cells_per_block=(2, 2),84 block_norm='L2-Hys',85 visualize=True,86 transform_sqrt=True87 )88 return features.astype(np.float32)89 90def extract_color_histogram(image, bins=32):91 """Extract color histogram features from HSV image"""92 hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)93 hist_h = np.histogram(hsv[:,:,0], bins=bins, range=(0, 180))[0]94 hist_s = np.histogram(hsv[:,:,1], bins=bins, range=(0, 256))[0]95 hist_v = np.histogram(hsv[:,:,2], bins=bins, range=(0, 256))[0]96 return np.concatenate([hist_h, hist_s, hist_v]).astype(np.float32)97 98def extract_lbp_features(image, radius=3, n_points=24):99 """Extract Local Binary Pattern features for texture"""100 from skimage.feature import local_binary_pattern101 gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)102 lbp = local_binary_pattern(gray, n_points, radius, method='uniform')103 hist, _ = np.histogram(lbp.ravel(), bins=n_points + 2)104 hist = hist.astype(np.float32)105 hist /= (hist.sum() + 1e-7) # Normalize106 return hist107 108def extract_features_from_image(image_data: bytes) -> np.ndarray:109 """110 Extract HOG, color histogram, and LBP features from uploaded image111 Same feature extraction as used in training112 """113 try:114 # Convert bytes to PIL Image115 image = Image.open(io.BytesIO(image_data))116 117 # Convert to numpy array and RGB118 img_array = np.array(image)119 if len(img_array.shape) == 2: # Grayscale120 img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)121 elif img_array.shape[2] == 4: # RGBA122 img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)123 124 # Resize to match training size125 img_resized = cv2.resize(img_array, (128, 128))126 127 # Extract HOG features128 hog_feat = extract_hog_features(img_resized)129 130 # Extract color histogram131 col_feat = extract_color_histogram(img_resized)132 133 # Extract LBP features134 lbp_feat = extract_lbp_features(img_resized)135 136 # Combine features137 combined_features = np.concatenate([hog_feat, col_feat, lbp_feat])138 139 return combined_features.reshape(1, -1)140 141 except Exception as e:142 raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")143 144@app.on_event("startup")145async def startup_event():146 """Load models on startup"""147 success = load_models()148 if not success:149 print("Warning: Could not load models. Please run training script first.")150 151@app.get("/", response_class=HTMLResponse)152async def root():153 """Serve the dashboard"""154 try:155 with open("dashboard/index.html", "r") as f:156 return HTMLResponse(content=f.read())157 except FileNotFoundError:158 return HTMLResponse(content="<h1>Cat vs Dog Classification</h1><p>Dashboard not found. Please check dashboard folder.</p>")159 160@app.get("/api/health")161async def health_check():162 """Health check endpoint"""163 return {164 "status": "healthy",165 "model_loaded": model is not None,166 "best_kernel": metadata.get('best_kernel') if metadata else None167 }168 169@app.post("/api/predict")170async def predict_image(file: UploadFile = File(...)):171 """Predict cat or dog from uploaded image"""172 if not model:173 raise HTTPException(status_code=503, detail="Model not loaded")174 175 if not scaler:176 raise HTTPException(status_code=503, detail="Scaler not loaded")177 178 if not label_encoder:179 raise HTTPException(status_code=503, detail="Label encoder not loaded")180 181 try:182 # Read image data183 image_data = await file.read()184 185 # Extract features186 features = extract_features_from_image(image_data)187 features_scaled = scaler.transform(features)188 189 # Extract individual features for display190 img_array = np.array(Image.open(io.BytesIO(image_data)))191 if len(img_array.shape) == 2: # Grayscale192 img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)193 elif img_array.shape[2] == 4: # RGBA194 img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)195 img_resized = cv2.resize(img_array, (128, 128))196 197 hog_feat = extract_hog_features(img_resized)198 col_feat = extract_color_histogram(img_resized)199 lbp_feat = extract_lbp_features(img_resized)200 201 # Make prediction202 prediction = model.predict(features_scaled)[0]203 class_id = int(prediction)204 class_name = label_encoder.inverse_transform([class_id])[0]205 206 # Get confidence if available207 confidence = None208 if hasattr(model, 'decision_function'):209 decision_values = model.decision_function(features_scaled)[0]210 exp_values = np.exp(decision_values - np.max(decision_values))211 probabilities = exp_values / np.sum(exp_values)212 confidence = float(np.max(probabilities))213 214 return {215 "prediction": {216 "class_id": class_id,217 "class_name": class_name,218 "confidence": confidence219 },220 "features": {221 "hog_size": len(hog_feat),222 "color_size": len(col_feat),223 "lbp_size": len(lbp_feat)224 }225 }226 227 except Exception as e:228 raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")229 230if __name__ == "__main__":231 import uvicorn232 uvicorn.run(app, host="0.0.0.0", port=7860)233