Modelo de detección de fallos 3D basado en Yolov5

Cómo usar

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)


### Combina con el modelo original

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


### Prepara una imagen

#### Del dataset original

from datasets import load_dataset

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

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

#### Desde local

from PIL import Image

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


### Inferencia y muestra la detección

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': 'parte', 'color': (255,0,0)},
{'name': 'espagueti', '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