Spaces:
Sleeping
Sleeping
File size: 1,020 Bytes
5bab6dc 5cb1a46 86001ee 5cb1a46 21011da 5cb1a46 299d49b ff9efc0 feca41c 5cb1a46 21011da 5cb1a46 ce3d5e4 4c6db5e 86001ee 299d49b 86001ee a1ef609 809b3d0 a1ef609 86001ee 809b3d0 3e53a2c 809b3d0 |
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 |
from ultralytics import YOLO
import asyncio
from fastapi import FastAPI, File, UploadFile, HTTPException
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.post("/predict")
async def get_prediction(file: UploadFile = File(...)):
file_name = file.filename
print(file_name)
with open(file_name, "wb") as file_object:
file_object.write(file.file.read())
# Predict with the model
results = model(file_name)
# View results
class_name = ""
for r in results:
print(names[int(r.probs.top1)])
class_name = names[int(r.probs.top1)]
return class_name |