Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| import numpy as np | |
| import cv2 | |
| from tensorflow.keras.models import load_model | |
| from tensorflow.keras.utils import to_categorical | |
| import os | |
| # FastAPI app setup | |
| app = FastAPI() | |
| # Load the pre-trained model | |
| model_path = 'sample.h5' | |
| if not os.path.exists(model_path): | |
| raise FileNotFoundError(f"Model file '{model_path}' not found. Ensure it is available in the working directory.") | |
| model = load_model(model_path) | |
| # Predefined image size for prediction | |
| size = 100 | |
| # FastAPI Model for input data | |
| class ImageData(BaseModel): | |
| image_path: str # Path to image file | |
| async def read_root(): | |
| """ | |
| Root endpoint for checking if the API is running. | |
| """ | |
| return {"message": "FastAPI app for pothole detection is running successfully!"} | |
| # Prediction Endpoint | |
| async def predict(data: ImageData): | |
| """ | |
| Predict whether the input image is a pothole or not. | |
| """ | |
| # Validate the image path | |
| if not os.path.exists(data.image_path): | |
| raise HTTPException(status_code=400, detail=f"Image file '{data.image_path}' not found.") | |
| try: | |
| # Load the image and preprocess | |
| image = cv2.imread(data.image_path, 0) | |
| if image is None: | |
| raise ValueError("Failed to read the image. Please check the file format and path.") | |
| image = cv2.resize(image, (size, size)) | |
| image = np.asarray(image).reshape(1, size, size, 1) | |
| # Make prediction | |
| prediction = model.predict(image) | |
| predicted_class = np.argmax(prediction, axis=1) | |
| # Return prediction result | |
| return {"predicted_class": int(predicted_class[0])} | |
| except Exception as e: | |
| # Handle unexpected errors | |
| raise HTTPException(status_code=500, detail=f"Error during prediction: {str(e)}") | |