Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,835 Bytes
ed1b7ea |
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 |
import typing
import ultralytics
YOLO_V10_MODELS = {
"nano": "jameslahm/yolov10n",
"small": "jameslahm/yolov10s",
"medium": "jameslahm/yolov10m",
"base": "jameslahm/yolov10b",
"large": "jameslahm/yolov10l",
"xlarge": "jameslahm/yolov10x",
}
class YOLOv10Plugin:
def __init__(
self,
yolo_model_name: (
str
| typing.Literal[
"nano",
"small",
"medium",
"base",
"large",
"xlarge",
]
) = "nano",
verbose: bool = True,
):
assert (
yolo_model_name in YOLO_V10_MODELS.keys()
), f"`yolo_model_name` should be either one of {list(YOLO_V10_MODELS.keys())}"
self.yolo_model_name = yolo_model_name
self.model = ultralytics.YOLOv10.from_pretrained(
YOLO_V10_MODELS[yolo_model_name]
)
self.verbose = verbose
if self.verbose:
print(f"YOLOv10Plugin::__init__::Model Name: {self.yolo_model_name}")
def detect(self, image):
results = self.model(image)
results = results[0].summary()
out = []
for result in results:
out.append(
{
"name": result["name"],
"class": result["class"],
"confidence": result["confidence"],
"box": [
int(result["box"]["x1"]),
int(result["box"]["y1"]),
int(result["box"]["x2"]),
int(result["box"]["y2"]),
],
}
)
return out
if __name__ == "__main__":
yolo = YOLOv10Plugin()
yolo.detect("https://ultralytics.com/images/zidane.jpg")
|