Spaces:
Sleeping
Sleeping
File size: 3,173 Bytes
bf132d0 7afb625 bf132d0 7afb625 bf132d0 7afb625 bf132d0 7afb625 bf132d0 7afb625 bf132d0 7afb625 bf132d0 91b12c0 bf132d0 7afb625 bf132d0 7afb625 b4a4875 7afb625 bf132d0 7afb625 bf132d0 7afb625 ff99f90 bf132d0 7afb625 bf132d0 7afb625 91b12c0 7afb625 | 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 | import io
import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from ultralytics import YOLO
app = FastAPI(
title="Malaria YOLO Object Detector",
description="API for detecting malaria parasites using YOLOv8 with bounding boxes.",
version="2.0.0"
)
# Load the newly trained YOLO model
MODEL_PATH = "best.pt"
try:
model = YOLO(MODEL_PATH)
print("✅ YOLO Model loaded successfully.")
except Exception as e:
print(f"❌ Error loading YOLO model: {e}")
model = None
# Note: We no longer need the preprocess_image function because YOLO handles
# all the resizing and normalization mathematically behind the scenes!
@app.post("/api/predict")
async def predict_cell(file: UploadFile = File(...)):
if model is None:
raise HTTPException(status_code=500, detail="Model is not loaded.")
try:
contents = await file.read()
nparr = np.frombuffer(contents, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if img is None:
raise ValueError("Invalid image file format.")
# Run YOLO inference
# Since we only trained for 25 epochs, we drop the confidence threshold to 0.1
# to ensure it aggressively highlights suspicious cells even if it isn't 100% sure yet.
results = model.predict(img, conf=0.1)
result = results[0]
detections = []
highest_parasite_conf = 0.0
parasite_detected = False
# Parse every bounding box found in the image
for box in result.boxes:
# Box coordinates
x1, y1, x2, y2 = box.xyxy[0].tolist()
# Confidence score
conf = float(box.conf[0])
# Class ID (0=parasite, 1=healthy, etc)
cls_id = int(box.cls[0])
label = result.names[cls_id]
if "parasite" in label.lower():
parasite_detected = True
if conf > highest_parasite_conf:
highest_parasite_conf = conf
detections.append({
"label": label,
"confidence": conf,
"box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2}
})
return JSONResponse(content={
"filename": file.filename,
"parasite_detected": parasite_detected,
"parasite_probability": f"{highest_parasite_conf * 100:.2f}%" if parasite_detected else "0.00%",
"prediction_label": "Parasite" if parasite_detected else "Healthy",
"detections": detections,
"image_width": img.shape[1],
"image_height": img.shape[0]
})
except ValueError as ve:
raise HTTPException(status_code=400, detail=str(ve))
except Exception as e:
raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")
# Mount the web UI
app.mount("/", StaticFiles(directory="static", html=True), name="static")
# To run locally: uvicorn app:app --reload
|