| 34 | |
| 35 | # 后处理函数 |
| 36 | def postprocess(results, output_dir, img_path, model_name): |
| 37 | # 检查输出目录 |
| 38 | if not os.path.exists(output_dir): |
| 39 | os.mkdir(output_dir) |
| 40 | |
| 41 | # 读取图像 |
| 42 | img = cv2.imread(img_path) |
| 43 | |
| 44 | # 计算MASK |
| 45 | mask = (results[0][0] > 0).astype('float32') |
| 46 | |
| 47 | # 缩放 |
| 48 | h, w = img.shape[:2] |
| 49 | mask = cv2.resize(mask, (w, h)) |
| 50 | |
| 51 | # 计算输出图像 |
| 52 | result = (img * mask[..., np.newaxis] + (1 - mask[..., np.newaxis]) * 255).astype(np.uint8) |
| 53 | |
| 54 | # 格式还原 |
| 55 | mask = (mask * 255).astype(np.uint8) |
| 56 | |
| 57 | # 可视化 |
| 58 | cv2.imwrite(os.path.join(output_dir, 'result_mask_%s.png' % model_name), mask) |
| 59 | cv2.imwrite(os.path.join(output_dir, 'result_%s.png' % model_name), result) |