Spaces:
Sleeping
Sleeping
| 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'} | |
| 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 |