Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, File, UploadFile, HTTPException | |
| from fastapi.responses import JSONResponse | |
| import tensorflow as tf | |
| from tensorflow.keras.preprocessing.image import img_to_array | |
| from PIL import Image | |
| import numpy as np | |
| import io | |
| # Initialize FastAPI app | |
| app = FastAPI(title="Cat vs Dog Classifier API") | |
| # Load the pre-trained model | |
| model = tf.keras.models.load_model('model.h5') | |
| # Define class labels | |
| class_names = ['Cat', 'Dog'] | |
| async def predict(file: UploadFile = File(...)): | |
| # Check if the uploaded file is an image | |
| if not file.content_type.startswith('image/'): | |
| raise HTTPException(status_code=400, detail="File must be an image") | |
| # Read and preprocess the image | |
| try: | |
| contents = await file.read() | |
| image = Image.open(io.BytesIO(contents)).convert('RGB') | |
| image = image.resize((224, 224)) # Resize to match model input | |
| image_array = img_to_array(image) / 255.0 # Rescale to [0, 1] | |
| image_array = np.expand_dims(image_array, axis=0) # Add batch dimension | |
| # Make prediction | |
| prediction = model.predict(image_array) | |
| predicted_class = class_names[int(prediction[0][0] > 0.5)] # Sigmoid threshold | |
| confidence = float(prediction[0][0]) if predicted_class == 'Dog' else float(1 - prediction[0][0]) | |
| return JSONResponse({ | |
| "predicted_class": predicted_class, | |
| "confidence": confidence | |
| }) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}") | |
| async def root(): | |
| return {"message": "Welcome to the Cat vs Dog Classifier API. Use POST /predict/ to classify an image."} |