Update app.py

app.py

@@ -21,19 +21,12 @@ templates = Jinja2Templates(directory="templates")
 
 def predict_yolo(image_path):
     # Load a model
-    model = YOLO('/usr/src/ultralytics/yolov8n.pt')  # pretrained YOLOv8n model
+    model = YOLO('ultralyticsplus/yolov8s')
 
     # Run batched inference on a list of images
     results = model(image_path)  # return a list of Results objects
 
-    # Process results list
-    for result in results:
-        boxes = result.boxes  # Boxes object for bbox outputs
-        # masks = result.masks  # Masks object for segmentation masks outputs
-        # keypoints = result.keypoints  # Keypoints object for pose outputs
-        # probs = result.probs  # Probs object for classification outputs
-    predictions = boxes.json()
-    return predictions
+    return results
 
 
 def draw_boxes(image, boxes):
@@ -43,29 +36,30 @@ def draw_boxes(image, boxes):
     return image
 
 
-@app.post("/uploadfile")
-async def create_upload_file(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents))
-    
-    # Save the image to a static directory
-    save_path = f"/tmp/{file.filename}"
-    image.save(save_path)
-    
-    # Perform YOLO prediction
-    predictions = predict_yolo(save_path)
-    
-    # Draw bounding boxes on the image
-    image_np = np.array(image)
-    image_with_boxes = draw_boxes(image_np, predictions)
-    
-    # Save the image with bounding boxes
-    image_with_boxes_path = f"/tmp/{file.filename.split('.')[0]}_with_boxes.jpg"
-    cv2.imwrite(image_with_boxes_path, cv2.cvtColor(image_with_boxes, cv2.COLOR_RGB2BGR))
-    
-    # Render the HTML with the image and bounding boxes
-    return templates.TemplateResponse("prediction.html", {"request": file, "image_path": image_with_boxes_path})
+@app.post("/", response_class=HTMLResponse)
+async def detect_yolo(request: Request, url: str = Form(...)
+    try:
+        # Download the image from the specified URL
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url)
+            response.raise_for_status()  # Raise an exception if there is an error in the request
+            content = response.content
+
+        image = Image.open(io.BytesIO(content))
+
+        results = predict_yolo(image)
+
+        render = render_result(model=model, image=image, result=results[0])
+
+        # Save the modified image to a byte sequence
+        image_byte_array = io.BytesIO()
+        image.save(image_byte_array, format="PNG")
+        
+        # Return the image as a response with content type "image/png"
+        return templates.TemplateResponse("result.html", {"request": request, "image": base64.b64encode(image_byte_array.getvalue()).decode()})
 
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
 @app.get("/test")
 async def read_root():
     return {"message": "TEST"}