-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
40 lines (31 loc) · 1.2 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from PIL import Image, ImageDraw
from ultralytics import YOLO
import torch
class Analyzer:
def __init__(self, model):
self.model = model
def process_image(self, image_path):
# Load image
img = Image.open(image_path)
# Run inference
with torch.no_grad():
output = self.model(img)
bounding_boxes = []
for result in output[0].boxes.xyxy.cpu().numpy():
x1, y1, x2, y2 = map(int, result[:4])
bounding_boxes.append([x1, y1, x2 - x1, y2 - y1])
# Save the image with bounding boxes
self.save_image_with_boxes(img, bounding_boxes, output_path='output.jpg')
def save_image_with_boxes(self, img, bounding_boxes, output_path):
# Convert image to RGB mode
img = img.convert('RGB')
# Draw bounding boxes on the image
draw = ImageDraw.Draw(img)
for box in bounding_boxes:
x, y, width, height = map(int, box)
draw.rectangle([x, y, x + width, y + height], outline="red", width=2)
# Save the image
img.save(output_path)
# Example usage for testing
analyzer = Analyzer(model=YOLO('porosity_model.pt'))
analyzer.process_image('example.jpeg')