diff --git a/examples/object_detector/yolo/yolov8/README.md b/examples/object_detector/yolo/yolov8/README.md new file mode 100644 index 0000000000..dcfe975c9c --- /dev/null +++ b/examples/object_detector/yolo/yolov8/README.md @@ -0,0 +1,64 @@ +# Object Detection using Ultralytics's pretrained YOLOv8(yolov8n) model. + + +Install `ultralytics` using +``` +python -m pip install -r requirements.txt +``` + +In this example, we are using the YOLOv8 Nano model from ultralytics.Downlaod the pretrained weights from [Ultralytics](https://docs.ultralytics.com/models/yolov8/#supported-modes) + +``` +wget https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt +``` + +We need a custom handler to load the YOLOv8n model. The default `initialize` function loads `.pt` file using `torch.jit.load`. This doesn't work for YOLOv8n model. Hence, we need a custom handler with an `initialize` method where we load the model using ultralytics. + +## Create a model archive file for Yolov8n model + +``` +torch-model-archiver --model-name yolov8n --version 1.0 --serialized-file yolov8n.pt --handler custom_handler.py +``` + +``` +mkdir model_store +mv yolov8n.mar model_store/. +``` + +## Start TorchServe and register the model + + +``` +torchserve --start --model-store model_store --ncs +curl -X POST "localhost:8081/models?model_name=yolov8n&url=yolov8n.mar&initial_workers=4&batch_size=2" +``` + +results in + +``` +{ + "status": "Model \"yolov8n\" Version: 1.0 registered with 4 initial workers" +} +``` + +## Run Inference + +Here we are counting the number of detected objects in the image. You can change the post-process method in the handler to return the bounding box coordinates + +``` +curl http://127.0.0.1:8080/predictions/yolov8n -T persons.jpg & curl http://127.0.0.1:8080/predictions/yolov8n -T bus.jpg +``` + +gives the output + +``` +{ + "person": 4, + "handbag": 3, + "bench": 3 +}{ + "person": 4, + "bus": 1, + "stop sign": 1 +} +``` diff --git a/examples/object_detector/yolo/yolov8/bus.jpg b/examples/object_detector/yolo/yolov8/bus.jpg new file mode 100644 index 0000000000..b43e311165 Binary files /dev/null and b/examples/object_detector/yolo/yolov8/bus.jpg differ diff --git a/examples/object_detector/yolo/yolov8/custom_handler.py b/examples/object_detector/yolo/yolov8/custom_handler.py new file mode 100644 index 0000000000..c91779c079 --- /dev/null +++ b/examples/object_detector/yolo/yolov8/custom_handler.py @@ -0,0 +1,79 @@ +import logging +import os +from collections import Counter + +import torch +from torchvision import transforms +from ultralytics import YOLO + +from ts.torch_handler.object_detector import ObjectDetector + +logger = logging.getLogger(__name__) + +try: + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +except ImportError as error: + XLA_AVAILABLE = False + + +class Yolov8Handler(ObjectDetector): + image_processing = transforms.Compose( + [transforms.Resize(640), transforms.CenterCrop(640), transforms.ToTensor()] + ) + + def __init__(self): + super(Yolov8Handler, self).__init__() + + def initialize(self, context): + # Set device type + if torch.cuda.is_available(): + self.device = torch.device("cuda") + elif XLA_AVAILABLE: + self.device = xm.xla_device() + else: + self.device = torch.device("cpu") + + # Load the model + properties = context.system_properties + self.manifest = context.manifest + model_dir = properties.get("model_dir") + self.model_pt_path = None + if "serializedFile" in self.manifest["model"]: + serialized_file = self.manifest["model"]["serializedFile"] + self.model_pt_path = os.path.join(model_dir, serialized_file) + self.model = self._load_torchscript_model(self.model_pt_path) + logger.debug("Model file %s loaded successfully", self.model_pt_path) + + self.initialized = True + + def _load_torchscript_model(self, model_pt_path): + """Loads the PyTorch model and returns the NN model object. + + Args: + model_pt_path (str): denotes the path of the model file. + + Returns: + (NN Model Object) : Loads the model object. + """ + # TODO: remove this method if https://github.com/pytorch/text/issues/1793 gets resolved + + model = YOLO(model_pt_path) + model.to(self.device) + return model + + def postprocess(self, res): + output = [] + for data in res: + classes = data.boxes.cls.tolist() + names = data.names + + # Map to class names + classes = map(lambda cls: names[int(cls)], classes) + + # Get a count of objects detected + result = Counter(classes) + output.append(dict(result)) + + return output diff --git a/examples/object_detector/yolo/yolov8/persons.jpg b/examples/object_detector/yolo/yolov8/persons.jpg new file mode 100644 index 0000000000..861d56a033 Binary files /dev/null and b/examples/object_detector/yolo/yolov8/persons.jpg differ diff --git a/examples/object_detector/yolo/yolov8/requirements.txt b/examples/object_detector/yolo/yolov8/requirements.txt new file mode 100644 index 0000000000..c5db96688c --- /dev/null +++ b/examples/object_detector/yolo/yolov8/requirements.txt @@ -0,0 +1 @@ +ultralytics>=8.0.144