-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
68 lines (50 loc) · 1.95 KB
/
main.py
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
from fastapi import FastAPI,HTTPException,UploadFile,File
import io
from io import BytesIO
from ultralytics import YOLO
import pandas as pd
import os
import numpy as np
import cv2
from fastapi.responses import RedirectResponse,StreamingResponse
import json
from Helper.helperFunc import add_BoundingBoxes,get_Images_from_Bytes,get_bytes_from_Images,get_model_predict,save_image
app=FastAPI(title="Detect Plastic")
@app.on_event("startup")
def startupevent():
if not os.path.exists("Prediction"):
os.makedirs("Prediction")
# if not os.path.exists("tmp"):
# os.makedirs("tmp")
print("Start Done")
@app.get('/')
async def redirect():
return RedirectResponse("/docs")
@app.post("/predict_to_get_plastics")
async def predict_to_json(file: UploadFile = File(...)):
model = YOLO("use saample or trained model")
image = get_Images_from_Bytes(file.file.read())
predictions = get_model_predict(image, model,flag=True)
data = json.loads(predictions)
info= {
"no_of_plastics": len(data),
"predictions": data
}
return info
@app.post("/predict_to_image")
def predict_to_image(file: UploadFile = File(...)):
model = YOLO("YOLO_Custom_v8m.pt")
img = get_Images_from_Bytes(file.file.read())
predictions = get_model_predict(img, model)
# print(predictions)
processed = add_BoundingBoxes(predictions, img)
return StreamingResponse(content=get_bytes_from_Images(processed) , media_type="image/jpeg")
@app.post("/predict_save_image")
async def predict_and_save(file:UploadFile=File(...)):
model=YOLO("use saample or trained model")
img=get_Images_from_Bytes(file.file.read())
predictions=get_model_predict(img,model)
bb_box=add_BoundingBoxes(img,predictions)
processed_image_bytes=get_bytes_from_Images(bb_box)
save_image(file_name=file.filename,image=bb_box,prediction=predictions)
return StreamingResponse(content=processed_image_bytes,media_type="image/jpeg")