safras_ml_api / main.py
Arafath10's picture
Update main.py
ce3d5e4 verified
raw
history blame
708 Bytes
from ultralytics import YOLO
import asyncio
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import requests
model = YOLO('best.pt') # load a custom model
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
names = {0: 'breaking-stage', 1: 'half-riping-stage', 2: 'overripe un-healthy', 3: 'ripe', 4: 'ripe_with_consumable_disease', 5: 'unripe'}
@app.get("/predict")
async def get_prediction():
# Predict with the model
results = model('test.jpg')
# View results
for r in results:
print(names[int(r.probs[top1])])
return r.probs