Spaces:
Sleeping
Sleeping
| import torch | |
| import cv2 | |
| import IPython | |
| from PIL import ImageColor | |
| from ultralytics import YOLO | |
| class ObjectDetection: | |
| def __init__(self, model_name='yolov3'): | |
| self.model_name = model_name | |
| self.model = self.load_model() | |
| self.classes = self.model.names | |
| self.device = 'cpu' | |
| def load_model(self): | |
| if self.model_name != 'yolov8': | |
| model = torch.hub.load(f"ultralytics/{self.model_name}", 'custom', path=f"./weights/{self.model_name}_best.pt", force_reload=True) | |
| else: | |
| model = YOLO(f"./weights/{self.model_name}_best.pt") | |
| return model | |
| def score_frame(self, frame): | |
| results = self.model(frame) | |
| labels, conf, coord = results.xyxyn[0][:, -1], results.xyxyn[0][:, -2], results.xyxyn[0][:, :-1] | |
| return labels, conf, coord | |
| def v8_score_frame(self, frame): | |
| results = self.model(frame) | |
| labels = [] | |
| confidences = [] | |
| coords = [] | |
| for result in results: | |
| boxes = result.boxes.cpu().numpy() | |
| label = boxes.cls | |
| conf = boxes.conf | |
| coord = boxes.xyxy | |
| labels.extend(label) | |
| confidences.extend(conf) | |
| coords.extend(coord) | |
| return labels, confidences, coords | |
| def get_coords(self, frame, row): | |
| if self.model_name != 'yolov8': | |
| x_shape, y_shape = frame.shape[1], frame.shape[0] | |
| return int(row[0] * x_shape), int(row[1] * y_shape), int(row[2] * x_shape), int(row[3] * y_shape) | |
| else: | |
| return int(row[0]), int(row[1]), int(row[2]), int(row[3]) | |
| def class_to_label(self, x): | |
| return self.classes[int(x)] | |
| def get_color(self, code): | |
| rgb = ImageColor.getcolor(code, "RGB") | |
| return rgb | |
| def plot_bboxes(self, results, frame, threshold, box_color, text_color): | |
| labels, conf, coord = results | |
| n = len(labels) | |
| frame = frame.copy() | |
| box_color = self.get_color(box_color) | |
| text_color = self.get_color(text_color) | |
| for i in range(n): | |
| row = coord[i] | |
| if conf[i] >= threshold: | |
| x1, y1, x2, y2 = self.get_coords(frame, row) | |
| class_name = self.class_to_label(labels[i]) | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), box_color, 2) | |
| cv2.putText(frame, f"{class_name} - {conf[i]*100:.2f}%", (x1, y1), cv2.FONT_HERSHEY_COMPLEX, 0.5, text_color) | |
| return frame | |