Create main.py

main.py

@@ -0,0 +1,30 @@
+from fastapi import FastAPI, File, UploadFile, Response
+from PIL import Image
+import torch
+import io
+
+app = FastAPI()
+
+# Load the YOLOv5 model
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
+
+@app.post("/detect/")
+async def detect(file: UploadFile = File(...)):
+    # Read image file
+    image_data = await file.read()
+    image = Image.open(io.BytesIO(image_data))
+
+    # Perform inference
+    results = model(image)
+
+    # Render the results on the image
+    results.render()
+
+    # Convert the image to bytes
+    img_bytes = io.BytesIO()
+    image_with_boxes = Image.fromarray(results.ims[0])
+    image_with_boxes.save(img_bytes, format='JPEG')
+    img_bytes.seek(0)
+
+    # Create a response with the image
+    return Response(content=img_bytes.getvalue(), media_type="image/jpeg")