3D Failures detection model based on Yolov5

This model was created using YOLOv5 in ultralytics Hub with the '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

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

model = torch.hub.load('Ultralytics/yolov5', 'custom', model_path, verbose = False)

Prepare an image

From the original dataset

from datasets import load_dataset

dataset = load_dataset('Javiai/failures-3D-print')

image = dataset["train"][0]["image"]

From local

from PIL import Image

image = Image.load("path/to/image")

Inference and show the detection

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