Views
No views yet
YOLOv5 in ultralytics Hub with the 'Javiai/failures-3D-print' dataset.1from huggingface_hub import hf_hub_download
2import torch
3
4repo_id = "Javiai/3dprintfails-yolo5vs"
5filename = "model_torch.pt"
6
7model_path = hf_hub_download(repo_id=repo_id, filename=filename)model = torch.hub.load('Ultralytics/yolov5', 'custom', model_path, verbose = False)1from datasets import load_dataset
2
3dataset = load_dataset('Javiai/failures-3D-print')
4
5image = dataset["train"][0]["image"]1from PIL import Image
2
3image = Image.load("path/to/image")
41from PIL import ImageDraw
2
3draw = ImageDraw.Draw(image)
4
5detections = model(image)
6
7categories = [
8 {'name': 'error', 'color': (0,0,255)},
9 {'name': 'extrusor', 'color': (0,255,0)},
10 {'name': 'part', 'color': (255,0,0)},
11 {'name': 'spaghetti', 'color': (0,0,255)}
12]
13
14for detection in detections.xyxy[0]:
15 x1, y1, x2, y2, p, category_id = detection
16 x1, y1, x2, y2, category_id = int(x1), int(y1), int(x2), int(y2), int(category_id)
17 draw.rectangle((x1, y1, x2, y2),
18 outline=categories[category_id]['color'],
19 width=1)
20 draw.text((x1, y1), categories[category_id]['name'],
21 categories[category_id]['color'])
22
23image
24