import math import os import time from collections import namedtuple import cv2 import numpy as np import torch from torchvision.transforms import ToTensor from DMPRUtils.model import DirectionalPointDetector from conf import config from utils.datasets import letterbox from utils.general import clip_coords from utils.torch_utils import select_device MarkingPoint = namedtuple('MarkingPoint', ['x', 'y', 'direction', 'shape']) def plot_points(image, pred_points, line_thickness=3): """Plot marking points on the image.""" if pred_points.size: tl = line_thickness or round(0.002 * (image.shape[0] + image.shape[1]) / 2) + 1 # line/font thickness tf = max(tl - 1, 1) # font thickness for conf, *point in pred_points: p0_x, p0_y = int(point[0]), int(point[1]) cos_val = math.cos(point[2]) sin_val = math.sin(point[2]) p1_x = int(p0_x + 20 * cos_val * tl) p1_y = int(p0_y + 20 * sin_val * tl) p2_x = int(p0_x - 10 * sin_val * tl) p2_y = int(p0_y + 10 * cos_val * tl) p3_x = int(p0_x + 10 * sin_val * tl) p3_y = int(p0_y - 10 * cos_val * tl) cv2.line(image, (p0_x, p0_y), (p1_x, p1_y), (0, 0, 255), thickness=tl) cv2.putText(image, str(float(conf)), (p0_x, p0_y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 0), thickness=tf) if point[3] > 0.5: cv2.line(image, (p0_x, p0_y), (p2_x, p2_y), (0, 0, 255), thickness=tl) else: cv2.line(image, (p2_x, p2_y), (p3_x, p3_y), (0, 0, 255), thickness=tf) def preprocess_image(image): """Preprocess numpy image to torch tensor.""" if image.shape[0] != 640 or image.shape[1] != 640: image = cv2.resize(image, (640, 640)) return torch.unsqueeze(ToTensor()(image), 0) def non_maximum_suppression(pred_points): """Perform non-maxmum suppression on marking points.""" suppressed = [False] * len(pred_points) for i in range(len(pred_points) - 1): for j in range(i + 1, len(pred_points)): i_x = pred_points[i][1].x i_y = pred_points[i][1].y j_x = pred_points[j][1].x j_y = pred_points[j][1].y # 0.0625 = 1 / 16 if abs(j_x - i_x) < 0.0625 and abs(j_y - i_y) < 0.0625: idx = i if pred_points[i][0] < pred_points[j][0] else j suppressed[idx] = True if any(suppressed): unsupres_pred_points = [] for i, supres in enumerate(suppressed): if not supres: unsupres_pred_points.append(pred_points[i]) return unsupres_pred_points return pred_points def get_predicted_points(prediction, thresh): """Get marking points from one predicted feature map.""" assert isinstance(prediction, torch.Tensor) predicted_points = [] prediction = prediction.detach().cpu().numpy() for i in range(prediction.shape[1]): for j in range(prediction.shape[2]): if prediction[0, i, j] >= thresh: xval = (j + prediction[2, i, j]) / prediction.shape[2] yval = (i + prediction[3, i, j]) / prediction.shape[1] # if not (config.BOUNDARY_THRESH <= xval <= 1-config.BOUNDARY_THRESH # and config.BOUNDARY_THRESH <= yval <= 1-config.BOUNDARY_THRESH): # continue cos_value = prediction[4, i, j] sin_value = prediction[5, i, j] direction = math.atan2(sin_value, cos_value) marking_point = MarkingPoint( xval, yval, direction, prediction[1, i, j]) predicted_points.append((prediction[0, i, j], marking_point)) return non_maximum_suppression(predicted_points) def get_predicted_points2(prediction, thresh): """Get marking points from one predicted feature map.""" assert isinstance(prediction, torch.Tensor) # predicted_points = [] # prediction = prediction.detach().cpu().numpy() # for i in range(prediction.shape[1]): # for j in range(prediction.shape[2]): # if prediction[0, i, j] >= thresh: # xval = (j + prediction[2, i, j]) / prediction.shape[2] # yval = (i + prediction[3, i, j]) / prediction.shape[1] # # if not (config.BOUNDARY_THRESH <= xval <= 1-config.BOUNDARY_THRESH # # and config.BOUNDARY_THRESH <= yval <= 1-config.BOUNDARY_THRESH): # # continue # cos_value = prediction[4, i, j] # sin_value = prediction[5, i, j] # direction = math.atan2(sin_value, cos_value) # marking_point = MarkingPoint( # xval, yval, direction, prediction[1, i, j]) # predicted_points.append((prediction[0, i, j], marking_point)) prediction = prediction.permute(1, 2, 0).contiguous() # prediction (20, 20, 6) height = prediction.shape[0] width = prediction.shape[1] j = torch.arange(prediction.shape[1], device=prediction.device).float().repeat(prediction.shape[0], 1).unsqueeze(dim=2) i = torch.arange(prediction.shape[0], device=prediction.device).float().view(prediction.shape[0], 1).repeat(1, prediction.shape[1]).unsqueeze(dim=2) prediction = torch.cat((prediction, j, i), dim=2) # 过滤小于thresh的置信度 prediction = prediction[prediction[..., 0] > thresh] prediction[..., 2] = (prediction[..., 2] + prediction[..., 6]) / width prediction[..., 3] = (prediction[..., 3] + prediction[..., 7]) / height direction = torch.atan2(prediction[..., 5], prediction[..., 4]) prediction = torch.stack((prediction[..., 0], prediction[..., 2], prediction[..., 3], direction, prediction[..., 1]), dim=1) # return non_maximum_suppression(predicted_points) return prediction def detect_marking_points(detector, image, thresh, device): """Given image read from opencv, return detected marking points.""" # t1 = time.time() # torch.cuda.synchronize(device) prediction = detector(preprocess_image(image).to(device)) # torch.cuda.synchronize(device) # t2 = time.time() # print(f'detector: {t2 - t1:.3f}s') return get_predicted_points2(prediction[0], thresh) def scale_coords2(img1_shape, coords, img0_shape, ratio_pad=None): # Rescale coords (xy) from img1_shape to img0_shape if ratio_pad is None: # calculate from img0_shape gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding else: gain = ratio_pad[0][0] pad = ratio_pad[1] # 百分比x，y转换为实际x，y height, width = img1_shape coords[:, 0] = torch.round(width * coords[:, 0] - 0.5) coords[:, 1] = torch.round(height * coords[:, 1] - 0.5) coords[:, 0] -= pad[0] # x padding coords[:, 1] -= pad[1] # y padding coords[:, :3] /= gain #恢复成原始图片尺寸 coords[:, 0].clamp_(0, img0_shape[1]) coords[:, 1].clamp_(0, img0_shape[0]) return coords def DMPR_process(img0, model, device, args): height, width, _ = img0.shape img, ratio, (dw, dh) = letterbox(img0, args.dmprimg_size, auto=False) det = detect_marking_points(model, img, args.dmpr_thresh, device) # if not pred: # return torch.tensor([]) # # 由list转为tensor # det = torch.tensor([[conf, *tup] for conf, tup in pred]) if len(det): det[:, 1:3] = scale_coords2(img.shape[:2], det[:, 1:3], img0.shape) # conf, x, y, θ, shape return det if __name__ == '__main__': impath = r'I:\zjc\weiting1\Images' file = 'DJI_0001_8.jpg' imgpath = os.path.join(impath, file) img0 = cv2.imread(imgpath) device_ = '0' device = select_device(device_) args = config.get_parser_for_inference().parse_args() model = DirectionalPointDetector(3, args.depth_factor, config.NUM_FEATURE_MAP_CHANNEL).to(device) weights = r"E:\pycharmProject\DMPR-PS\weights\dp_detector_499.pth" model.load_state_dict(torch.load(weights)) det = DMPR_process(img0, model, device, args) plot_points(img0, det) cv2.imwrite(file, img0, [int(cv2.IMWRITE_JPEG_QUALITY), 100])