|
--- |
|
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 |
|
|
|
 |
