result_fusion_pass.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the needed packages\n",
    "# from result_fusion import Fusion\n",
    "from fusion.fusion_pass import Fusion_PASS\n",
    "import os\n",
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "from tools.iou_perimg import SegmentationMetrics\n",
    "import tqdm\n",
    "from config import config_pass as config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# folder pathes\n",
    "mask_folder = config.mask_folder\n",
    "mask_suffix = config.mask_suffix\n",
    "mask_color_folder = config.mask_color_folder\n",
    "mask_color_suffix = config.mask_color_suffix\n",
    "\n",
    "# the path to the model prediction\n",
    "segmentation_folder = config.segmentation_folder\n",
    "segmentation_suffix = config.segmentation_suffix\n",
    "segmentation_suffix_noimg = config.segmentation_suffix_noimg\n",
    "# \n",
    "confidence_folder = config.confidence_folder\n",
    "confidence_suffix = config.confidence_suffix\n",
    "entropy_folder = config.entropy_folder\n",
    "entropy_suffix = config.entropy_suffix\n",
    "\n",
    "# the path to the original image\n",
    "image_folder = config.image_folder\n",
    "image_suffix = config.image_suffix\n",
    "# the path to the ground truth\n",
    "gt_folder = config.gt_folder\n",
    "gt_suffix = config.gt_suffix\n",
    "no_gt = config.no_gt\n",
    "# the path to the output folder\n",
    "output_folder = config.output_folder\n",
    "output_suffix = config.output_suffix\n",
    "\n",
    "#num of classes\n",
    "num_classes = config.num_classes\n",
    "\n",
    "### fusion parameters\n",
    "# the fusion mode\n",
    "fusion_mode = config.fusion_mode\n",
    "road_assumption = config.road_assumption\n",
    "road_center_rect = config.road_center_rect\n",
    "get_sam_mode= config.get_sam_mode\n",
    "use_sgml = config.use_sgml\n",
    "sam_alpha = config.sam_alpha\n",
    "large_classes = config.large_classes  # [0, 1, 2, 8, 10, 13]\n",
    "small_classes = config.small_classes  # [3, 4, 5, 6, 7, 9, 11, 12, 14, 15, 16, 17, 18]\n",
    "sam_classes = config.sam_classes  # 11 classes, 5, 6, 7, \n",
    "shrink_num = config.shrink_num\n",
    "\n",
    "if 'BiSeNet-uda' in config.segmentation_folder:\n",
    "    confidence_threshold = config.confidence_threshold_tufl\n",
    "    entropy_ratio = config.entropy_ratio_tufl\n",
    "elif 'MIC' in config.segmentation_folder:\n",
    "    confidence_threshold = config.confidence_threshold_mic\n",
    "    entropy_ratio = config.entropy_ratio_mic\n",
    "else:\n",
    "    confidence_threshold = config.confidence_threshold_daformer\n",
    "    entropy_ratio = config.entropy_ratio_daformer\n",
    "\n",
    "### display parameters\n",
    "display_size = config.display_size\n",
    "mix_ratio = config.mix_ratio\n",
    "resize_ratio = config.resize_ratio\n",
    "\n",
    "### save params\n",
    "save_mix_result = config.save_mix_result\n",
    "save_sam_result = config.save_sam_result\n",
    "save_all_fusion = config.save_all_fusion  # [True, True, True, False, False]\n",
    "\n",
    "save_majority_process = config.save_majority_process\n",
    "save_sgml_process = config.save_sgml_process\n",
    "save_f1_process = config.save_f1_process\n",
    "save_f2_process = config.save_f2_process\n",
    "save_f3_process = config.save_f3_process\n",
    "\n",
    "### time setting\n",
    "time_process = config.time_process\n",
    "time_filename = config.time_filename\n",
    "\n",
    "### sample index setting\n",
    "debug_num = config.debug_num # 2975\n",
    "begin_index = config.begin_index # 0\n",
    "\n",
    "fusion = Fusion_PASS(mask_folder, \n",
    "                    mask_suffix,\n",
    "                    mask_color_folder,\n",
    "                    mask_color_suffix,\n",
    "                    segmentation_folder, \n",
    "                    segmentation_suffix,\n",
    "                    segmentation_suffix_noimg,\n",
    "                    confidence_folder,\n",
    "                    confidence_suffix,\n",
    "                    entropy_folder,\n",
    "                    entropy_suffix,\n",
    "                    image_folder,\n",
    "                    image_suffix,\n",
    "                    gt_folder, \n",
    "                    gt_suffix,\n",
    "                    output_folder,\n",
    "                    output_suffix,\n",
    "                    num_classes,\n",
    "                    fusion_mode,\n",
    "                    road_assumption,\n",
    "                    road_center_rect,\n",
    "                    get_sam_mode,\n",
    "                    sam_alpha,\n",
    "                    large_classes,\n",
    "                    small_classes,\n",
    "                    sam_classes,\n",
    "                    shrink_num,\n",
    "                    display_size,\n",
    "                    mix_ratio,\n",
    "                    resize_ratio,\n",
    "                    time_process,\n",
    "                    time_filename,\n",
    "                    save_sgml_process,\n",
    "                    save_majority_process,\n",
    "                    save_f1_process,\n",
    "                    save_f2_process,\n",
    "                    save_f3_process\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_range = list(range(begin_index, begin_index + debug_num))\n",
    "iou_cal = SegmentationMetrics(num_classes=num_classes)\n",
    "\n",
    "if save_all_fusion:\n",
    "    f1_output_folder = os.path.join(output_folder, 'fusion1_trainid')\n",
    "    f1_color_output_folder = os.path.join(output_folder, 'fusion1_color')\n",
    "    f2_output_folder = os.path.join(output_folder, 'fusion2_trainid')\n",
    "    f2_color_output_folder = os.path.join(output_folder, 'fusion2_color')\n",
    "    f3_output_folder = os.path.join(output_folder, 'fusion3_trainid')\n",
    "    f3_color_output_folder = os.path.join(output_folder, 'fusion3_color')\n",
    "    # f4_output_folder = os.path.join(output_folder, 'fusion4_trainid')\n",
    "    # f4_color_output_folder = os.path.join(output_folder, 'fusion4_color')\n",
    "    # f5_output_folder = os.path.join(output_folder, 'fusion5_trainid')\n",
    "    fusion.check_and_make(f1_output_folder)\n",
    "    fusion.check_and_make(f1_color_output_folder)\n",
    "    fusion.check_and_make(f2_output_folder)\n",
    "    fusion.check_and_make(f2_color_output_folder)\n",
    "    fusion.check_and_make(f3_output_folder)\n",
    "    fusion.check_and_make(f3_color_output_folder)\n",
    "    # fusion.check_and_make(f4_output_folder)\n",
    "    # fusion.check_and_make(f4_color_output_folder)\n",
    "    # fusion.check_and_make(f5_output_folder)\n",
    "\n",
    "\n",
    "if save_sam_result:\n",
    "    sam_majority_output_folder = os.path.join(output_folder, 'sam_majority_trainid')\n",
    "    sam_majority_color_output_folder = os.path.join(output_folder, 'sam_majority_color')\n",
    "    sam_sgml_output_folder = os.path.join(output_folder, 'sam_sgml_trainid')\n",
    "    sam_sgml_color_output_folder = os.path.join(output_folder, 'sam_sgml_color')\n",
    "    fusion.check_and_make(sam_majority_output_folder)\n",
    "    fusion.check_and_make(sam_majority_color_output_folder)\n",
    "    fusion.check_and_make(sam_sgml_output_folder)\n",
    "    fusion.check_and_make(sam_sgml_color_output_folder)\n",
    "\n",
    "bar = tqdm.tqdm(total=debug_num)\n",
    "for i in index_range:\n",
    "    # get the image name\n",
    "    image_name = fusion.image_names[i]  # 1\n",
    "    \n",
    "    # get color mask\n",
    "    mask_color_path = os.path.join(fusion.mask_color_folder, image_name + fusion.mask_color_suffix)\n",
    "    mask_color = cv2.imread(mask_color_path, 1) # [h, w, 3]\n",
    "    \n",
    "    # get the prediction\n",
    "    # aachen_000000_000019_leftImg8bittrainID.png\n",
    "    prediction_path = os.path.join(fusion.segmentation_folder, image_name + fusion.segmentation_suffix)\n",
    "    # print(',',fusion.segmentation_folder, image_name + fusion.segmentation_suffix)\n",
    "    uda_pred = cv2.imread(prediction_path, 0) # [h, w], 1 rgb channel\n",
    "    uda_pred_color = fusion.color_segmentation(uda_pred)\n",
    "    \n",
    "    # get the confidence map\n",
    "    confidence_path = os.path.join(fusion.confidence_folder, image_name + fusion.confidence_suffix)\n",
    "    pred_confidence = np.load(confidence_path, allow_pickle=True) # [h, w]\n",
    "    \n",
    "    # get the entropy map\n",
    "    entropy_path = os.path.join(fusion.entropy_folder, image_name + fusion.entropy_suffix)\n",
    "    pred_entropy = np.load(entropy_path, allow_pickle=True) # [h, w]\n",
    "    \n",
    "    # get the ground truth\n",
    "    if not no_gt:\n",
    "        gt_path = os.path.join(fusion.gt_folder, image_name + fusion.gt_suffix)\n",
    "        gt = cv2.imread(gt_path, 0) # [h, w]\n",
    "        gt_color = fusion.color_segmentation(gt)\n",
    "    else:\n",
    "        gt_color = np.zeros((uda_pred.shape[0], uda_pred.shape[1], 3), dtype=np.uint8)\n",
    "        mask_color = np.zeros((uda_pred.shape[0], uda_pred.shape[1], 3), dtype=np.uint8)\n",
    "\n",
    "    # get the original image\n",
    "    original_image = cv2.imread(os.path.join(fusion.image_folder, image_name + fusion.image_suffix))\n",
    "\n",
    "    # get the confidence and entropy map\n",
    "    pred_confidence = pred_confidence.astype(np.float32)\n",
    "    confidence_map = fusion.visualize_numpy(pred_confidence)\n",
    "    confidence_mask, confidence_img = fusion.vis_np_higher_thres(pred_confidence, original_image, confidence_threshold)\n",
    "    \n",
    "    pred_entropy = pred_entropy.astype(np.float32)\n",
    "    entropy_map = fusion.visualize_numpy(pred_entropy)\n",
    "    entropy_threshold = np.percentile(pred_entropy, entropy_ratio)\n",
    "    entropy_mask, entropy_img = fusion.vis_np_lower_thres(pred_entropy, original_image, entropy_threshold)\n",
    "\n",
    "    # get the sam segmentation result using the mask\n",
    "    sam_pred_sgml, sam_pred_majority = fusion.get_sam_pred(image_name, uda_pred, confidence_mask, entropy_mask)  # [h,w]\n",
    "    sam_color_sgml = fusion.color_segmentation(sam_pred_sgml)  # [h,w,3]\n",
    "    \n",
    "    if use_sgml:\n",
    "        sam_pred = sam_pred_sgml\n",
    "    else:\n",
    "        sam_pred = sam_pred_majority\n",
    "    \n",
    "    # save the sam result\n",
    "    image_filename = image_name + fusion.mask_suffix\n",
    "    if save_sam_result:\n",
    "        sam_majority_color = fusion.color_segmentation(sam_pred_majority) # [h,w,3]\n",
    "        cv2.imwrite(os.path.join(sam_majority_output_folder, image_filename), sam_pred_majority)\n",
    "        cv2.imwrite(os.path.join(sam_majority_color_output_folder, image_filename), sam_majority_color)\n",
    "        cv2.imwrite(os.path.join(sam_sgml_output_folder, image_filename), sam_pred_sgml)\n",
    "        cv2.imwrite(os.path.join(sam_sgml_color_output_folder, image_filename), sam_color_sgml)\n",
    "\n",
    "    # initialize the fusion color results and error images with black image\n",
    "    black_img = np.zeros((original_image.shape[0], original_image.shape[1], 3), dtype=np.uint8)\n",
    "    fusion_color_bg_1, fusion_color_bg_2, fusion_color_bg_3 = black_img, black_img, black_img\n",
    "    error_0, error_1, error_2, error_3 = black_img, black_img, black_img, black_img\n",
    "    \n",
    "    # initialize fusion trainid with 255 one-channel image\n",
    "    ignore_lb = np.ones((original_image.shape[0], original_image.shape[1]), dtype=np.uint8) * 255\n",
    "    fusion_trainid_bg_1, fusion_trainid_bg_2, fusion_trainid_bg_3 = ignore_lb, ignore_lb, ignore_lb\n",
    "    \n",
    "    # get fusion result from 1, 2, 3, 4, 5\n",
    "    fusion_trainid_bg_1, fusion_color_bg_1 = \\\n",
    "        fusion.fusion_mode_1(uda_pred=uda_pred, sam_pred=sam_pred, image_name=image_name)\n",
    "    fusion_trainid_bg_2, fusion_color_bg_2 = \\\n",
    "        fusion.fusion_mode_2(uda_pred=uda_pred, sam_pred=sam_pred, image_name=image_name)\n",
    "    fusion_trainid_bg_3, fusion_color_bg_3 = \\\n",
    "        fusion.fusion_mode_3(uda_pred=uda_pred, sam_pred=sam_pred, fusion_trainid=fusion_trainid_bg_1, confidence_mask=confidence_mask, entropy_mask=entropy_mask, image_name=image_name)\n",
    "    # fusion_trainid_bg_4, fusion_color_bg_4 = \\\n",
    "    #     fusion.fusion_mode_4(uda_pred=uda_pred, sam_pred=sam_pred, \\\n",
    "    #                          fusion_trainid=fusion_trainid_bg_3, confidence_mask=confidence_mask)\n",
    "    # fusion_trainid_bg_5, fusion_color_bg_5 = \\\n",
    "    #     fusion.fusion_mode_5(uda_pred=uda_pred, sam_pred=sam_pred, \\\n",
    "    #                          fusion_trainid=fusion_trainid_bg_3, entropy_mask=entropy_mask)\n",
    "    \n",
    "    miou_0, miou_1, miou_2, miou_3, miou_4, miou_5 = -1, -1, -1, -1, -1, -1\n",
    "    ious_0, ious_1, ious_2, ious_3, ious_4, ious_5 = [], [], [], [], [], []\n",
    "    if not no_gt:\n",
    "        miou_0, ious_0 = iou_cal.calculate_miou(uda_pred, gt)\n",
    "        miou_1, ious_1 = iou_cal.calculate_miou(fusion_trainid_bg_1, gt)\n",
    "        miou_2, ious_2 = iou_cal.calculate_miou(fusion_trainid_bg_2, gt)\n",
    "        miou_3, ious_3 = iou_cal.calculate_miou(fusion_trainid_bg_3, gt)\n",
    "        # miou_4, ious_4 = iou_cal.calculate_miou(fusion_trainid_bg_4, gt)\n",
    "        # miou_5, ious_5 = iou_cal.calculate_miou(fusion_trainid_bg_5, gt)\n",
    "    \n",
    "        error_0 = fusion.get_error_image(uda_pred, gt, uda_pred_color)\n",
    "        error_1 = fusion.get_error_image(fusion_trainid_bg_1, gt, fusion_color_bg_1)\n",
    "        error_2 = fusion.get_error_image(fusion_trainid_bg_2, gt, fusion_color_bg_2)\n",
    "        error_3 = fusion.get_error_image(fusion_trainid_bg_3, gt, fusion_color_bg_3)\n",
    "        # error_4 = fusion.get_error_image(fusion_trainid_bg_4, gt, fusion_color_bg_4)\n",
    "        # error_5 = fusion.get_error_image(fusion_trainid_bg_5, gt, fusion_color_bg_5)\n",
    "    \n",
    "    # display the results\n",
    "    fusion.dis_imgs_horizontal(\n",
    "        [original_image, gt_color, mask_color, \\\n",
    "         sam_color_sgml, uda_pred_color, error_0, \\\n",
    "         fusion_color_bg_1, fusion_color_bg_2, fusion_color_bg_3, \\\n",
    "         error_1, error_2, error_3, \\\n",
    "         confidence_map, entropy_map, confidence_img], \\\n",
    "        '{}'.format(image_name), \\\n",
    "        [(miou_0, ious_0), (miou_1, ious_1), (miou_2, ious_2), \\\n",
    "         (miou_3, ious_3)], \\\n",
    "        [confidence_threshold, entropy_threshold])\n",
    "    \n",
    "    # save the mious and ious\n",
    "    if not no_gt:\n",
    "        miou_values = [miou_0, miou_1, miou_2, miou_3]\n",
    "        ious_values = [ious_0, ious_1, ious_2, ious_3]\n",
    "        fusion.save_ious(miou_values, ious_values, '{}'.format(image_name))\n",
    "    \n",
    "    # save all fusion results\n",
    "    if save_all_fusion:\n",
    "        image_filename = image_name + fusion.output_suffix\n",
    "        cv2.imwrite(os.path.join(f1_output_folder, image_filename), fusion_trainid_bg_1)\n",
    "        cv2.imwrite(os.path.join(f1_color_output_folder, image_filename), fusion_color_bg_1)\n",
    "        cv2.imwrite(os.path.join(f2_output_folder, image_filename), fusion_trainid_bg_2)\n",
    "        cv2.imwrite(os.path.join(f2_color_output_folder, image_filename), fusion_color_bg_2)\n",
    "        cv2.imwrite(os.path.join(f3_output_folder, image_filename), fusion_trainid_bg_3)\n",
    "        cv2.imwrite(os.path.join(f3_color_output_folder, image_filename), fusion_color_bg_3)\n",
    "        # cv2.imwrite(os.path.join(f4_output_folder, image_filename), fusion_trainid_bg_4)\n",
    "        # cv2.imwrite(os.path.join(f4_color_output_folder, image_filename), fusion_color_bg_4)\n",
    "        # cv2.imwrite(os.path.join(f5_output_folder, image_filename), fusion_trainid_bg_5)\n",
    "    \n",
    "    # save mix results\n",
    "    if save_mix_result:\n",
    "        # get the sam mixed color image using the fusion.mix_ratio\n",
    "        sam_mixed_color = cv2.addWeighted(original_image, fusion.mix_ratio, sam_color_sgml, 1 - fusion.mix_ratio, 0)\n",
    "        if fusion.resize_ratio != 1:\n",
    "            new_h = int(sam_mixed_color.shape[0] * fusion.resize_ratio)\n",
    "            new_w = int(sam_mixed_color.shape[1] * fusion.resize_ratio)\n",
    "            sam_mixed_color = cv2.resize(sam_mixed_color, (new_w, new_h), interpolation=cv2.INTER_NEAREST)\n",
    "        \n",
    "        # save the sam mask in trainid and color to the output folder\n",
    "        cv2.imwrite(os.path.join(fusion.output_folder, 'trainID', image_filename), sam_pred)\n",
    "        # cv2.imwrite(os.path.join(fusion.output_folder, 'color', image_filename), fusion_color)\n",
    "        cv2.imwrite(os.path.join(fusion.output_folder, 'mixed', image_filename), sam_mixed_color)\n",
    "        \n",
    "        # make the fusion results to list for easy use    \n",
    "        fusion_trainid_bgs = [fusion_trainid_bg_1, fusion_trainid_bg_2, fusion_trainid_bg_3]\n",
    "        fusion_color_bgs = [fusion_color_bg_1, fusion_color_bg_2, fusion_color_bg_3]\n",
    "        \n",
    "        mode = fusion.fusion_mode - 1\n",
    "        if mode in range(len(fusion_trainid_bgs)):\n",
    "            fusion_trainid_bg, fusion_color_bg = fusion_trainid_bgs[mode], fusion_color_bgs[mode]\n",
    "        else:\n",
    "            raise NotImplementedError(\"This fusion mode has not been implemented yet.\")\n",
    "    \n",
    "        #save the fusion mask in trainid and color to the output folder\n",
    "        mixed_color_bg = cv2.addWeighted(original_image, fusion.mix_ratio, fusion_color_bg, 1 - fusion.mix_ratio, 0)\n",
    "        if fusion.resize_ratio != 1:\n",
    "            mixed_color_bg = cv2.resize(mixed_color_bg, (int(mixed_color_bg.shape[1] * fusion.resize_ratio), int(mixed_color_bg.shape[0] * fusion.resize_ratio)), interpolation=cv2.INTER_NEAREST)\n",
    "        cv2.imwrite(os.path.join(fusion.output_folder, 'trainID_bg', image_filename), fusion_trainid_bg)\n",
    "        # cv2.imwrite(os.path.join(fusion.output_folder, 'color_bg', image_filename), fusion_color_bg)\n",
    "        cv2.imwrite(os.path.join(fusion.output_folder, 'mixed_bg', image_filename), mixed_color_bg)\n",
    "\n",
    "    bar.update(1)\n",
    "bar.close()\n",
    "if time_process:\n",
    "    fusion.show_time_process()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sam",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}