diff --git a/anylabeling/configs/autolabel_yolov5l.yaml b/anylabeling/configs/autolabel_yolov5l.yaml new file mode 100644 index 0000000..316bf82 --- /dev/null +++ b/anylabeling/configs/autolabel_yolov5l.yaml @@ -0,0 +1,90 @@ +type: yolov5 +name: yolov5l +display_name: YOLOv5l Ultralytics +model_path: anylabeling_assets/models/yolov5/yolov5l.onnx +input_width: 640 +input_height: 640 +score_threshold: 0.5 +nms_threshold: 0.45 +confidence_threshold: 0.45 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/autolabel_yolov5n.yaml b/anylabeling/configs/autolabel_yolov5n.yaml new file mode 100644 index 0000000..a6f8bf4 --- /dev/null +++ b/anylabeling/configs/autolabel_yolov5n.yaml @@ -0,0 +1,90 @@ +type: yolov5 +name: yolov5n +display_name: YOLOv5n Ultralytics +model_path: anylabeling_assets/models/yolov5/yolov5n.onnx +input_width: 640 +input_height: 640 +score_threshold: 0.5 +nms_threshold: 0.45 +confidence_threshold: 0.45 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/autolabel_yolov5x.yaml b/anylabeling/configs/autolabel_yolov5x.yaml new file mode 100644 index 0000000..8867083 --- /dev/null +++ b/anylabeling/configs/autolabel_yolov5x.yaml @@ -0,0 +1,90 @@ +type: yolov5 +name: yolov5x +display_name: YOLOv5x Ultralytics +model_path: anylabeling_assets/models/yolov5/yolov5x.onnx +input_width: 640 +input_height: 640 +score_threshold: 0.5 +nms_threshold: 0.45 +confidence_threshold: 0.45 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/autolabel_yolov8l.yaml b/anylabeling/configs/autolabel_yolov8l.yaml new file mode 100644 index 0000000..2ad8977 --- /dev/null +++ b/anylabeling/configs/autolabel_yolov8l.yaml @@ -0,0 +1,90 @@ +type: yolov8 +name: yolov8l +display_name: YOLOv8l Ultralytics +model_path: anylabeling_assets/models/yolov8/yolov8l.onnx +input_width: 640 +input_height: 640 +score_threshold: 0.5 +nms_threshold: 0.45 +confidence_threshold: 0.45 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/autolabel_yolov8m.yaml b/anylabeling/configs/autolabel_yolov8m.yaml new file mode 100644 index 0000000..a62081e --- /dev/null +++ b/anylabeling/configs/autolabel_yolov8m.yaml @@ -0,0 +1,90 @@ +type: yolov8 +name: yolov8m +display_name: YOLOv8m Ultralytics +model_path: anylabeling_assets/models/yolov8/yolov8m.onnx +input_width: 640 +input_height: 640 +score_threshold: 0.5 +nms_threshold: 0.45 +confidence_threshold: 0.45 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/autolabel_yolov8n.yaml b/anylabeling/configs/autolabel_yolov8n.yaml new file mode 100644 index 0000000..1d11655 --- /dev/null +++ b/anylabeling/configs/autolabel_yolov8n.yaml @@ -0,0 +1,90 @@ +type: yolov8 +name: yolov8n +display_name: YOLOv8n Ultralytics +model_path: anylabeling_assets/models/yolov8/yolov8n.onnx +input_width: 640 +input_height: 640 +score_threshold: 0.5 +nms_threshold: 0.45 +confidence_threshold: 0.45 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/autolabel_yolov8s.yaml b/anylabeling/configs/autolabel_yolov8s.yaml new file mode 100644 index 0000000..e11a4e5 --- /dev/null +++ b/anylabeling/configs/autolabel_yolov8s.yaml @@ -0,0 +1,90 @@ +type: yolov8 +name: yolov8s +display_name: YOLOv8s Ultralytics +model_path: anylabeling_assets/models/yolov8/yolov8s.onnx +input_width: 640 +input_height: 640 +score_threshold: 0.5 +nms_threshold: 0.45 +confidence_threshold: 0.45 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/configs/autolabel_yolov8x.yaml b/anylabeling/configs/autolabel_yolov8x.yaml new file mode 100644 index 0000000..251d96a --- /dev/null +++ b/anylabeling/configs/autolabel_yolov8x.yaml @@ -0,0 +1,90 @@ +type: yolov8 +name: yolov8x +display_name: YOLOv8x Ultralytics +model_path: anylabeling_assets/models/yolov8/yolov8x.onnx +input_width: 640 +input_height: 640 +score_threshold: 0.5 +nms_threshold: 0.45 +confidence_threshold: 0.45 +classes: + - person + - bicycle + - car + - motorcycle + - airplane + - bus + - train + - truck + - boat + - traffic light + - fire hydrant + - stop sign + - parking meter + - bench + - bird + - cat + - dog + - horse + - sheep + - cow + - elephant + - bear + - zebra + - giraffe + - backpack + - umbrella + - handbag + - tie + - suitcase + - frisbee + - skis + - snowboard + - sports ball + - kite + - baseball bat + - baseball glove + - skateboard + - surfboard + - tennis racket + - bottle + - wine glass + - cup + - fork + - knife + - spoon + - bowl + - banana + - apple + - sandwich + - orange + - broccoli + - carrot + - hot dog + - pizza + - donut + - cake + - chair + - couch + - potted plant + - bed + - dining table + - toilet + - tv + - laptop + - mouse + - remote + - keyboard + - cell phone + - microwave + - oven + - toaster + - sink + - refrigerator + - book + - clock + - vase + - scissors + - teddy bear + - hair drier + - toothbrush diff --git a/anylabeling/services/auto_labeling/model.py b/anylabeling/services/auto_labeling/model.py index ea08fee..2b015ab 100644 --- a/anylabeling/services/auto_labeling/model.py +++ b/anylabeling/services/auto_labeling/model.py @@ -11,7 +11,7 @@ class Model: BASE_DOWNLOAD_URL = ( - "https://github.com/vietanhdev/anylabeling-assets/raw/main/" + "https://github.com/vietanhdev/anylabeling-assets/releases/download/v0.0.1/" ) class Meta: @@ -57,7 +57,14 @@ def get_model_abs_path(self, model_path): " while..." ) relative_path = model_path.replace("anylabeling_assets/", "") - download_url = self.BASE_DOWNLOAD_URL + relative_path + + if relative_path.startswith("models/segment_anything/"): + download_url = self.BASE_DOWNLOAD_URL + relative_path.replace("models/segment_anything/","") + elif relative_path.startswith("models/yolov5/"): + download_url = self.BASE_DOWNLOAD_URL + relative_path.replace("models/yolov5/","") + else: + download_url = self.BASE_DOWNLOAD_URL + relative_path.replace("models/yolov8/","") + home_dir = os.path.expanduser("~") model_abs_path = os.path.abspath( os.path.join(home_dir, "anylabeling_data", relative_path) diff --git a/anylabeling/services/auto_labeling/model_manager.py b/anylabeling/services/auto_labeling/model_manager.py index 4133fde..04c9df0 100644 --- a/anylabeling/services/auto_labeling/model_manager.py +++ b/anylabeling/services/auto_labeling/model_manager.py @@ -21,8 +21,16 @@ class ModelManager(QObject): model_configs = { "segment_anything_vit_b": "autolabel_segment_anything.yaml", + "yolov5n": "autolabel_yolov5n.yaml", "yolov5s": "autolabel_yolov5s.yaml", "yolov5m": "autolabel_yolov5m.yaml", + "yolov5l": "autolabel_yolov5l.yaml", + "yolov5x": "autolabel_yolov5x.yaml", + "yolov8n": "autolabel_yolov8n.yaml", + "yolov8s": "autolabel_yolov8s.yaml", + "yolov8m": "autolabel_yolov8m.yaml", + "yolov8l": "autolabel_yolov8l.yaml", + "yolov8x": "autolabel_yolov8x.yaml", } def __init__(self): @@ -104,17 +112,23 @@ def _load_model(self, model_name): self.auto_segmentation_model_unselected.emit() model_info = copy.deepcopy(self.model_infos[model_name]) - if model_info["type"] == "yolov5": from .yolov5 import YOLOv5 - + model_info["model"] = YOLOv5( model_info, on_message=self.new_model_status.emit ) self.auto_segmentation_model_unselected.emit() + elif model_info["type"] == "yolov8": + from .yolov8 import YOLOv8 + + model_info["model"] = YOLOv8( + model_info, on_message=self.new_model_status.emit + ) + self.auto_segmentation_model_unselected.emit() elif model_info["type"] == "segment_anything": from .segment_anything import SegmentAnything - + model_info["model"] = SegmentAnything( model_info, on_message=self.new_model_status.emit ) @@ -150,4 +164,4 @@ def predict_shapes(self, image): auto_labeling_result = self.loaded_model_info["model"].predict_shapes( image ) - self.new_auto_labeling_result.emit(auto_labeling_result) + self.new_auto_labeling_result.emit(auto_labeling_result) \ No newline at end of file diff --git a/anylabeling/services/auto_labeling/segment_anything.py b/anylabeling/services/auto_labeling/segment_anything.py index 0b65adc..a6f63e5 100644 --- a/anylabeling/services/auto_labeling/segment_anything.py +++ b/anylabeling/services/auto_labeling/segment_anything.py @@ -54,11 +54,13 @@ def __init__(self, config_path, on_message) -> None: raise Exception(f"Decoder not found: {decoder_model_abs_path}") # Load models + cuda = True if onnxruntime.get_device() == 'GPU' else False + providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider'] self.encoder_session = onnxruntime.InferenceSession( - encoder_model_abs_path + encoder_model_abs_path, providers=providers ) self.decoder_session = onnxruntime.InferenceSession( - decoder_model_abs_path + decoder_model_abs_path, providers=providers ) # Mark for auto labeling diff --git a/anylabeling/services/auto_labeling/yolov8.py b/anylabeling/services/auto_labeling/yolov8.py new file mode 100644 index 0000000..bb26e29 --- /dev/null +++ b/anylabeling/services/auto_labeling/yolov8.py @@ -0,0 +1,172 @@ +import logging +import os + +import cv2 +import numpy as np +from PyQt5 import QtCore + +from anylabeling.views.labeling.shape import Shape +from anylabeling.views.labeling.utils.opencv import qt_img_to_cv_img + +from .model import Model +from .types import AutoLabelingResult + + +class YOLOv8(Model): + """Object detection model using YOLOv8""" + + class Meta: + required_config_names = [ + "type", + "name", + "display_name", + "model_path", + "input_width", + "input_height", + "score_threshold", + "nms_threshold", + "confidence_threshold", + "classes", + ] + buttons = ["button_run"] + + def __init__(self, model_config, on_message) -> None: + # Run the parent class's init method + super().__init__(model_config, on_message) + + model_abs_path = self.get_model_abs_path(self.config["model_path"]) + if not os.path.isfile(model_abs_path): + raise Exception(f"Model not found: {model_abs_path}") + + self.net = cv2.dnn.readNet(model_abs_path) + self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA) + self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA) + self.classes = self.config["classes"] + + def pre_process(self, input_image, net): + """ + Pre-process the input image before feeding it to the network. + """ + # Create a 4D blob from a frame. + blob = cv2.dnn.blobFromImage( + input_image, + 1 / 255, + (self.config["input_width"], self.config["input_height"]), + [0, 0, 0], + 1, + crop=False, + ) + + # Sets the input to the network. + net.setInput(blob) + + # Runs the forward with blob. + outputs = net.forward() + outputs = np.array([cv2.transpose(outputs[0])]) + + return outputs + + def post_process(self, input_image, outputs): + """ + Post-process the network's output, to get the bounding boxes and + their confidence scores. + """ + # Lists to hold respective values while unwrapping. + class_ids = [] + confidences = [] + boxes = [] + + # Rows. + rows = outputs.shape[1] + + image_height, image_width = input_image.shape[:2] + + # Resizing factor. + x_factor = image_width / self.config["input_width"] + y_factor = image_height / self.config["input_height"] + + # Iterate through 8400 rows. + for r in range(rows): + row = outputs[0][r] + classes_scores = row[4:] + + # Get the index of max class score and confidence. + _, confidence, _, (_, class_id) = cv2.minMaxLoc(classes_scores) + + # Discard confidence lower than threshold + if confidence >= self.config["confidence_threshold"]: + confidences.append(confidence) + class_ids.append(class_id) + + cx, cy, w, h = row[0], row[1], row[2], row[3] + + left = int((cx - w / 2) * x_factor) + top = int((cy - h / 2) * y_factor) + width = int(w * x_factor) + height = int(h * y_factor) + + box = np.array([left, top, width, height]) + boxes.append(box) + + # Perform non maximum suppression to eliminate redundant + # overlapping boxes with lower confidences. + indices = cv2.dnn.NMSBoxes( + boxes, + confidences, + self.config["confidence_threshold"], + self.config["nms_threshold"], + ) + + output_boxes = [] + for i in indices: + box = boxes[i] + left = box[0] + top = box[1] + width = box[2] + height = box[3] + label = self.classes[class_ids[i]] + score = confidences[i] + + output_box = { + "x1": left, + "y1": top, + "x2": left + width, + "y2": top + height, + "label": label, + "score": score, + } + + output_boxes.append(output_box) + + return output_boxes + + def predict_shapes(self, image): + """ + Predict shapes from image + """ + + if image is None: + return [] + + try: + image = qt_img_to_cv_img(image) + except Exception as e: # noqa + logging.warning("Could not inference model") + logging.warning(e) + return [] + + detections = self.pre_process(image, self.net) + boxes = self.post_process(image, detections) + shapes = [] + + for box in boxes: + shape = Shape(label=box["label"], shape_type="rectangle", flags={}) + shape.add_point(QtCore.QPointF(box["x1"], box["y1"])) + shape.add_point(QtCore.QPointF(box["x2"], box["y2"])) + shapes.append(shape) + + result = AutoLabelingResult(shapes, replace=True) + return result + + def unload(self): + del self.net diff --git a/pyproject.toml b/pyproject.toml index c5ba840..8d14fe4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,3 +14,8 @@ exclude = ''' | dist )/ ''' + +[build-system] +requires = ["setuptools", "wheel", "onnxruntime-gpu"] +build-backend = "setuptools.build_meta" + diff --git a/requirements-dev.txt b/requirements-dev.txt index 15e5816..40b2fb3 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,6 +5,7 @@ github2pypi==1.0.0 pre-commit==2.20.0 onnx==1.13.1 onnxruntime==1.14.1 +onnxruntime-gpu==1.14.1 qimage2ndarray==1.10.0 build twine diff --git a/requirements-macos.txt b/requirements-macos.txt index 49b3df3..4df6ca9 100644 --- a/requirements-macos.txt +++ b/requirements-macos.txt @@ -1,4 +1,4 @@ -opencv-contrib-python-headless==4.5.5.62 +opencv-contrib-python-headless==4.7.0.68 # PyQt5==5.15.6 # Use Miniconda/Anaconda: conda install -c conda-forge pyqt imgviz==1.5.0 natsort==8.1.0 @@ -8,4 +8,5 @@ PyYAML==6.0 colorama==0.4.5 onnx==1.13.1 onnxruntime==1.14.1 +onnxruntime-gpu==1.14.1 qimage2ndarray==1.10.0 diff --git a/requirements.txt b/requirements.txt index cf900bf..c9929da 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -opencv-contrib-python-headless==4.6.0.66 +opencv-contrib-python-headless==4.7.0.68 PyQt5==5.15.7 imgviz==1.5.0 natsort==8.1.0 @@ -8,4 +8,5 @@ PyYAML==6.0 colorama==0.4.5 onnx==1.13.1 onnxruntime==1.14.1 +onnxruntime-gpu==1.14.1 qimage2ndarray==1.10.0