Spaces:
Sleeping
Sleeping
| 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, | |
| } | |
| 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 | |
| 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 | |