Spaces:
Sleeping
Sleeping
File size: 2,285 Bytes
ab0369b 56157b6 ab0369b c2a39aa ab0369b 56157b6 ab0369b ba17a5a ab0369b 98ccaaa ab0369b ba17a5a ab0369b ba17a5a ab0369b ba17a5a ab0369b ba17a5a ab0369b ba17a5a 56157b6 ba17a5a ab0369b 56157b6 ab0369b 56157b6 |
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 67 68 69 70 71 72 73 74 75 76 |
import gradio as gr
import io
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import json
# Load your custom YOLO model
model = YOLO("fentanyl_oft.pt")
def detect_keypoints(image):
"""
Run YOLO inference and return keypoints data
"""
try:
# Convert PIL Image to numpy array
if isinstance(image, Image.Image):
image_np = np.array(image)
image_cv2 = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
else:
image_cv2 = image
# Run inference
results = model.predict(
source=image_cv2,
conf=0.05,
iou=0.7,
max_det=1,
imgsz=1440,
device='cpu',
verbose=False
)
keypoints_data = []
if results and len(results) > 0:
result = results[0]
if result.keypoints is not None:
kpts = result.keypoints.xy.cpu().numpy()
conf = result.keypoints.conf.cpu().numpy()
for i in range(kpts.shape[1]):
if i < len(kpts[0]):
x, y = kpts[0][i]
confidence = conf[0][i] if i < len(conf[0]) else 0.0
keypoints_data.append({
"id": i,
"x": float(x),
"y": float(y),
"confidence": float(confidence)
})
return {
"success": True,
"keypoints": keypoints_data,
"image_width": image_cv2.shape[1],
"image_height": image_cv2.shape[0],
"num_keypoints": len(keypoints_data)
}
except Exception as e:
return {"success": False, "error": str(e)}
# Create Gradio interface with API access enabled
iface = gr.Interface(
fn=detect_keypoints,
inputs=gr.Image(type="pil"),
outputs=gr.JSON(),
title="YOLO Keypoint Detection",
description="Upload an image to detect keypoints using custom YOLO model",
api_name="predict" # This enables API access at /api/predict
)
# Launch with API enabled
if __name__ == "__main__":
iface.launch(share=False) |