metadata
license: apache-2.0
datasets:
- Javiai/failures-3D-print
language:
- en
pipeline_tag: object-detection
#3D Failures detection model based on Yolov5
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
