--- license: apache-2.0 datasets: - Javiai/failures-3D-print language: - en pipeline_tag: object-detection --- # 3D Failures detection model based on Yolov5 This model was created using ```YOLOv5``` in ultralytics Hub with the [```'Javiai/failures-3D-print'```](https://huggingface.co/datasets/Javiai/failures-3D-print) dataset. The idea is detect some failures in a 3D printing process. This model detect the part that is been printing, the extrusor, some errors and if there is a spaghetti error type ## How to use ### Download the model ```python from huggingface_hub import hf_hub_download import torch repo_id = "Javiai/3dprintfails-yolo5vs" filenam = "model_torch.pt" model_path = hf_hub_download(repo_id=repo_id, filename=filename) ``` ### Combine with the original model ```python model = torch.hub.load('Ultralytics/yolov5', 'custom', model_path, verbose = False) ``` ### Prepare an image #### From the original dataset ```python from datasets import load_dataset dataset = load_dataset('Javiai/failures-3D-print') image = dataset["train"][0]["image"] ``` #### From local ```python from PIL import Image image = Image.load("path/to/image") ``` ### Inference and show the detection ```python from PIL import ImageDraw draw = ImageDraw.Draw(image) detections = model(image) categories = [ {'name': 'error', 'color': (0,0,255)}, {'name': 'extrusor', 'color': (0,255,0)}, {'name': 'part', 'color': (255,0,0)}, {'name': 'spaghetti', 'color': (0,0,255)} ] for detection in detections.xyxy[0]: x1, y1, x2, y2, p, category_id = detection x1, y1, x2, y2, category_id = int(x1), int(y1), int(x2), int(y2), int(category_id) draw.rectangle((x1, y1, x2, y2), outline=categories[category_id]['color'], width=1) draw.text((x1, y1), categories[category_id]['name'], categories[category_id]['color']) image ``` ## Example image ![image/png](https://cdn-uploads.huggingface.co/production/uploads/63c9c08a5fdc575773c7549b/3ZSkBvN0o8sSpQjGxdwJx.png)