import os import numpy as np import torch import torch.nn as nn from PIL import ImageDraw, ImageFont, Image from nets.model import YoloBody from utils.utils import (cvtColor, get_anchors, get_classes, preprocess_input, resize_image, show_config) from utils.utils_bbox import DecodeBox, DecodeBoxNP class YOLO(object): _defaults = { "model_path" : 'final_experiments/best_epoch_weights.pth', "classes_path" : 'model_data/rtts_classes.txt', "anchors_path" : 'model_data/yolo_anchors.txt', "anchors_mask" : [[6, 7, 8], [3, 4, 5], [0, 1, 2]], "input_shape" : [640, 640], "confidence" : 0.5, "nms_iou" : 0.3, "letterbox_image" : True, "cuda" : False, } @classmethod def get_defaults(cls, n): if n in cls._defaults: return cls._defaults[n] else: return "Unrecognized attribute name '" + n + "'" def __init__(self, **kwargs): self.__dict__.update(self._defaults) for name, value in kwargs.items(): setattr(self, name, value) self._defaults[name] = value self.class_names, self.num_classes = get_classes(self.classes_path) self.anchors, self.num_anchors = get_anchors(self.anchors_path) self.bbox_util = DecodeBox(self.anchors, self.num_classes, (self.input_shape[0], self.input_shape[1]), self.anchors_mask) self.colors = [(255, 209, 227), (126, 161, 255), (91, 188, 255), (255, 127, 62), (255, 250, 183)] self.generate() show_config(**self._defaults) def generate(self, onnx=False): self.net = YoloBody(self.anchors_mask, self.num_classes) self.device = torch.device('cuda' if self.cuda and torch.cuda.is_available() else 'cpu') try: state_dict = torch.load(self.model_path, map_location=self.device, weights_only=True) except TypeError: state_dict = torch.load(self.model_path, map_location=self.device) self.net.load_state_dict(state_dict) self.net = self.net.fuse().eval().to(self.device) print('{} model, and classes loaded.'.format(self.model_path)) if not onnx: if self.cuda and self.device.type == 'cuda': self.net = nn.DataParallel(self.net) self.net = self.net.cuda() def detect_image(self, image, crop = False, count = False): image_shape = np.array(np.shape(image)[0:2]) image = cvtColor(image) image_data = resize_image(image, (self.input_shape[1], self.input_shape[0]), self.letterbox_image) image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0) with torch.no_grad(): images = torch.from_numpy(image_data).to(self.device) outputs = self.net(images) outputs = self.bbox_util.decode_box(outputs) results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou) if results[0] is None: return image top_label = np.array(results[0][:, 6], dtype = 'int32') top_conf = results[0][:, 4] * results[0][:, 5] top_boxes = results[0][:, :4] font = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32')) thickness = int(max((image.size[0] + image.size[1]) // np.mean(self.input_shape), 1)) if count: print("top_label:", top_label) classes_nums = np.zeros([self.num_classes]) for i in range(self.num_classes): num = np.sum(top_label == i) if num > 0: print(self.class_names[i], " : ", num) classes_nums[i] = num print("classes_nums:", classes_nums) if crop: for i, c in list(enumerate(top_boxes)): top, left, bottom, right = top_boxes[i] top = max(0, np.floor(top).astype('int32')) left = max(0, np.floor(left).astype('int32')) bottom = min(image.size[1], np.floor(bottom).astype('int32')) right = min(image.size[0], np.floor(right).astype('int32')) dir_save_path = "img_crop" if not os.path.exists(dir_save_path): os.makedirs(dir_save_path) crop_image = image.crop([left, top, right, bottom]) crop_image.save(os.path.join(dir_save_path, "crop_" + str(i) + ".png"), quality=95, subsampling=0) print("save crop_" + str(i) + ".png to " + dir_save_path) for i, c in list(enumerate(top_label)): predicted_class = self.class_names[int(c)] box = top_boxes[i] score = top_conf[i] top, left, bottom, right = box top = max(0, np.floor(top).astype('int32')) left = max(0, np.floor(left).astype('int32')) bottom = min(image.size[1], np.floor(bottom).astype('int32')) right = min(image.size[0], np.floor(right).astype('int32')) label = '{} {:.2f}'.format(predicted_class, score) draw = ImageDraw.Draw(image) label_size = self._get_text_size(draw, label, font) label = label.encode('utf-8') print(label, top, left, bottom, right) if top - label_size[1] >= 0: text_origin = np.array([left, top - label_size[1]]) else: text_origin = np.array([left, top + 1]) for i in range(thickness): draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[c]) draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[c]) draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font) del draw return image @staticmethod def _get_text_size(draw, text, font): if hasattr(draw, 'textbbox'): left, top, right, bottom = draw.textbbox((0, 0), text, font=font) return right - left, bottom - top return draw.textsize(text, font) def get_map_txt(self, image_id, image, class_names, map_out_path): f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"), "w", encoding='utf-8') image_shape = np.array(np.shape(image)[0:2]) image = cvtColor(image) image_data = resize_image(image, (self.input_shape[1], self.input_shape[0]), self.letterbox_image) image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0) with torch.no_grad(): images = torch.from_numpy(image_data) if self.cuda: images = images.cuda() outputs = self.net(images) outputs = self.bbox_util.decode_box(outputs) results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou) if results[0] is None: return top_label = np.array(results[0][:, 6], dtype = 'int32') top_conf = results[0][:, 4] * results[0][:, 5] top_boxes = results[0][:, :4] for i, c in list(enumerate(top_label)): predicted_class = self.class_names[int(c)] box = top_boxes[i] score = str(top_conf[i]) top, left, bottom, right = box if predicted_class not in class_names: continue f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom)))) f.close() return