def seal_detect(img, model, device): boxes = detect.detect(img, model, device) boxes = detect.adjust_ratio(boxes, ratio_w, ratio_h) orig_vertices = [] theta = 0 if boxes is not None and boxes.size: for box in boxes: box = np.array(box[:8]) orig_vertices.append(box) theta += crop_mode.find_min_rect_angle(box) orig_vertices = np.array(orig_vertices) theta /= len(boxes) tmp_img, vertices = crop_mode.rotate_allimg(img, orig_vertices, -theta / math.pi * 180) dict_centers = {} for i, vertice in enumerate(vertices): avg_x = int(crop_mode.averagenum(vertice[::2])) avg_y = int(crop_mode.averagenum(vertice[1::2])) dict_centers[str(avg_x) + ',' + str(avg_y)] = i centers = crop_mode.sort_centers(dict_centers, 1) xcenters = [] for center in centers: xcenters.append([center]) shape = [] for i, xcenter in enumerate(xcenters): for center in xcenter: anno = {} anno['box'] = orig_vertices[int(center[1])] anno['class'] = 'seal' shape.append(anno) return shape, boxes
rotate_img = img.rotate(180, Image.BILINEAR) else: rotate_img = img detect_model.load_state_dict(torch.load(configs.detect_model_one)) detect_model.eval() rotate_img = rotate_img.convert("RGB") w, h = rotate_img.size ratio_w = 512 / w ratio_h = 512 / h img_tmp = rotate_img.resize((512, 512)) shape, seal_boxes = seal_detect(img_tmp, seal_model, device) boxes = detect.detect(img_tmp, detect_model, device) boxes = detect.adjust_ratio(boxes, ratio_w, ratio_h) plot_img = detect.plot_boxes(rotate_img, boxes) plot_img = detect.plot_boxes(plot_img, seal_boxes) plot_img.save('detected.jpg') print('detection result saved as: detected.jpg') orig_vertices = [] theta = 0 for box in boxes: box = np.array(box[:8]) orig_vertices.append(box) theta += crop_mode.find_min_rect_angle(box) orig_vertices = np.array(orig_vertices)