breedclassification / result.py
kshitiz14's picture
commit 1234567890abcdef1234567890abcdef12345678
3897b7c
Raw History Blame Contribute Delete
2.91 kB
import os
import uvicorn
import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse
from tensorflow.keras.models import load_model
app = FastAPI()
# Read allowed origins from env var (comma-separated)
ALLOWED_ORIGINS = [
"ALLOWED_ORIGINS",
"http://localhost:3000,"
"https://sih-em37.vercel.app"
]
app.add_middleware(
CORSMiddleware,
allow_origins=ALLOWED_ORIGINS,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Load model
model = load_model("my_image_model.h5")
class_indices = {
0: "Gir_cow",
1: "Murrah_buffalo",
2: "Red_Sindhi_cow",
3: "Sahiwal_cow",
4: "Tharparkar_cow",
5: "amritmahal_cow",
6: "banni_buffalo",
7: "bhadwari_buffalo",
8: "dharwadi_buffalo",
9: "jafarabadi_buffalo",
}
# Minimum confidence required
CONFIDENCE_THRESHOLD = 0.70
def preprocess_frame(frame, target_size=(224, 224)):
img = cv2.resize(frame, target_size)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.astype("float32") / 255.0
img = np.expand_dims(img, axis=0)
return img
@app.get("/")
async def index():
html_content = """
<html>
<body>
<h2>Upload Image for Prediction</h2>
<form action="/predict" enctype="multipart/form-data" method="post">
<input name="image" type="file">
<input type="submit" value="Upload">
</form>
</body>
</html>
"""
return HTMLResponse(content=html_content)
@app.post("/predict")
async def predict(image: UploadFile = File(...)):
try:
contents = await image.read()
np_arr = np.frombuffer(contents, np.uint8)
frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
if frame is None:
return JSONResponse(
content={"error": "Invalid image format"},
status_code=400
)
processed = preprocess_frame(frame)
preds = model.predict(processed, verbose=0)
class_id = int(np.argmax(preds, axis=1)[0])
confidence = float(np.max(preds))
# Reject low-confidence predictions
if confidence < CONFIDENCE_THRESHOLD:
return {
"label": "No cattle detected",
"confidence": round(confidence * 100, 2),
"detected": False
}
label = class_indices.get(class_id, f"Class_{class_id}")
return {
"label": label,
"confidence": round(confidence * 100, 2),
"detected": True
}
except Exception as e:
return JSONResponse(
content={"error": str(e)},
status_code=500
)
if __name__ == "__main__":
port = int(os.getenv("PORT", 7860))
uvicorn.run(app, host="0.0.0.0", port=port)