eval.py

import os
import cv2
import json
import numpy as np
import argparse
import torch
from mobile_sam import sam_model_registry, SamPredictor
import pycocotools.mask as mask_util
from tqdm import tqdm
import logging
import pycocotools.mask as mask_utils
from eval_tools import calculate_pa_iou, inference_image, init_model, overlay_mask_on_image, clean_checkpoint_path, calculate_metrics, inference_image_fps
from Tools_metrics.performance_process import log_performance_metrics, process_performance_metrics

os.environ['MODEL_MODE'] = "test"
os.environ['INFERENCE_MODE'] = "test"
def main():

    parser = argparse.ArgumentParser()
    # parser.add_argument("--model_type", type=str, default="tiny_msam", help="model type", choices=["vit_t","tiny_msam", "vit_h", "rep_sam", "micro_sam"])
    parser.add_argument("--model_type", type=str, default="vit_t", help="model type", choices=["vit_t","tiny_msam", "vit_h"])
    # parser.add_argument("--checkpoint_path", type=str, default="/data2/wuxinrui/Distill-SAM/trained_models/Img_Encoder_T_vit_t_S_tiny_msam/temp_copy/stroke_v8.pth", help="path to the checkpoint")
    parser.add_argument("--checkpoint_path", type=str, default="/data2/wuxinrui/Distill-SAM/trained_models/Img_Encoder_T_vit_t_S_vit_t/temp/last.pth", help="path to the checkpoint")
    # parser.add_argument("--checkpoint_path", type=str, default="/data2/wuxinrui/Distill-SAM/weights/mobile_sam.pt", help="path to the checkpoint")
    # parser.add_argument("--checkpoint_path", type=str, default="//data2/wuxinrui/Distill-SAM/trained_models/Img_Encoder_T_vit_t_S_micro_sam/temp/micro_sam_stage10.pth", help="path to the checkpoint")
    # parser.add_argument("--checkpoint_path", type=str, default="/data2/wuxinrui/Distill-SAM/weights/sam_vit_h_4b8939.pth", help="path to the checkpoint")
    parser.add_argument("--test_img_path", type=str, default="/data2/wuxinrui/Datasets/IMC_1000/MIMC_FINAL/seen/test_list", help="the test image path")
    parser.add_argument("--label_path", type=str, default="/data2/wuxinrui/Datasets/IMC_1000/jsons_for_salient_instance_segmentation/test_1_prompts.json", help="the test json path")
    # parser.add_argument("--test_img_path", type=str, default="/data2/wuxinrui/Datasets/IMC_1000/MIMC_FINAL/seen/train_list", help="the test image path")
    # parser.add_argument("--label_path", type=str, default="/data2/wuxinrui/Datasets/IMC_1000/MIMC_FINAL/train-taxonomic_cleaned.json", help="the test json path")
    parser.add_argument("--label_num", type=int, default=5, help="the num of points // more prior than label_path")
    parser.add_argument("--image_size", type=int, default=1024, help="image size")
    parser.add_argument("--visualize_mask_path", default="/data2/wuxinrui/Distill-SAM/modified_mobilesam_results")
    parser.add_argument("--ori_SAM", default=False)
    parser.add_argument("--sample_num", type=int, default=None, help="sample_num")
    
    args = parser.parse_args()
    
    checkpoint_name = os.path.basename(args.checkpoint_path)
    log_filename = f"/data2/wuxinrui/Distill-SAM/baseline_eval/eval_{checkpoint_name}.log"
    logging.basicConfig(
        filename=log_filename,
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"

    visualize_mask_path = args.visualize_mask_path
    
    if os.path.exists(visualize_mask_path):
        for file in os.listdir(visualize_mask_path):
            os.remove(os.path.join(visualize_mask_path, file))
    else:
        os.makedirs(visualize_mask_path)
        
    logging.info("Using device: " + device)
    logging.info(f"Devices num is {torch.cuda.device_count()}")
    logging.info(f"args: {args}")


    if args.label_num:
        if args.label_num == 1:
            label_path = "/data2/wuxinrui/Datasets/IMC_1000/jsons_for_salient_instance_segmentation/test_1_prompts.json"
        elif args.label_num == 3:
            label_path = "/data2/wuxinrui/Datasets/IMC_1000/jsons_for_salient_instance_segmentation/test_3_prompts.json"
        elif args.label_num == 5:
            label_path = "/data2/wuxinrui/Datasets/IMC_1000/jsons_for_salient_instance_segmentation/test_5_prompts.json"
    else:
        label_path = args.label_path
        
    with open(label_path, 'r') as f:
        json_data = json.load(f)
        annotations = json_data['annotations']

    img_files = [f for f in os.listdir(args.test_img_path) if f.endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif'))]
    
    bbox_pred_masks = []
    point_pred_masks = []
    
    bbox_metric_dicts = None
    point_metric_dicts = None
    temp_checkpoint_path = clean_checkpoint_path(args.checkpoint_path)
    # mask_generator, mask_predictor = init_model(model_type=args.model_type, sam_checkpoint=temp_checkpoint_path, device=device, generator=False, predictor=True, ori_SAM=True)
    mask_generator, mask_predictor = init_model(model_type=args.model_type, sam_checkpoint=temp_checkpoint_path, device=device, generator=False, predictor=True, ori_SAM=args.ori_SAM)


    logging.info("#" * 50 + f"INFERENCE_MODE = {os.environ['INFERENCE_MODE']}" + "#" * 50)
    logging.info("#" * 50 + f"label_path = {label_path}" + "#" * 50)
    

    sample_num = args.sample_num
    accumulated_metrics = None
    for img_file in tqdm(img_files[:sample_num], miniters=50):
        try:

            img_path = os.path.join(args.test_img_path, img_file)
            ext_name = img_file.split('.')[-1]
            point_prompt = True
            # results_dict = inference_image_fps(img_path, annotations, mask_generator, mask_predictor, device=device, bbox_prompt=True, point_prompt=point_prompt)
            results_dict = inference_image_fps(img_path, annotations, mask_generator, mask_predictor, device=device, bbox_prompt=True, point_prompt=point_prompt)
            
            predictor_results = results_dict['predictor_results']
            image_masks = results_dict['image_masks']
            
            single_masks = results_dict['single_masks']
            current_metrics, accumulated_metrics, avg_metrics = process_performance_metrics(
            results_dict['performance_metrics'], 
            accumulated_metrics, single_masks=single_masks
            )

            bbox_pred_masks = predictor_results['bbox']['masks']
            point_pred_masks =  predictor_results['point']['masks'] if point_prompt else bbox_pred_masks
            # point_pred_masks = predictor_results['point']['masks']
            
            output_path_bbox = os.path.join(visualize_mask_path, img_file.replace(ext_name, "_bbox_pred.jpg"))
            output_path_point = os.path.join(visualize_mask_path, img_file.replace(ext_name, "_point_pred.jpg"))
            
            overlay_mask_on_image(bbox_pred_masks, output_path=output_path_bbox, image_path=img_path)
            overlay_mask_on_image(point_pred_masks, output_path=output_path_point, image_path=img_path)
            
            
            bbox_metric_dict = calculate_metrics(bbox_pred_masks, image_masks)
            if bbox_metric_dicts is None:
                bbox_metric_dicts = {}
                for metric_name, metric_list in bbox_metric_dict.items():
                    bbox_metric_dicts[metric_name] = []
            for metric_name, metric_list in bbox_metric_dict.items():
                bbox_metric_dicts[metric_name].extend(metric_list[i] for i in range(len(metric_list)))


            point_metric_dict = calculate_metrics(point_pred_masks, image_masks)
            if point_metric_dicts is None:
                point_metric_dicts = {}
                for metric_name, metric_list in point_metric_dict.items():
                    point_metric_dicts[metric_name] = []
            for metric_name, metric_list in point_metric_dict.items():
                point_metric_dicts[metric_name].extend(metric_list[i] for i in range(len(metric_list)))
        
        except Exception as e:
            print(f"Error processing image {img_file}: {e}")
            continue

        # pa_list, iou_list = calculate_pa_iou(bbox_pred_masks, image_masks)
        # bbox_iou_list.extend(iou_list[i] for i in range(len(iou_list)))
        # bbox_pa_list.extend(pa_list[i] for i in range(len(pa_list)))
        
        # pa_list, iou_list = calculate_pa_iou(point_pred_masks, image_masks)
        # point_iou_list.extend(iou_list[i] for i in range(len(iou_list)))
        # point_pa_list.extend(pa_list[i] for i in range(len(pa_list)))
        
    
    # bbox_avg_iou = np.mean(bbox_iou_list)
    # bbox_avg_pa = np.mean(bbox_pa_list)
    # point_avg_iou = np.mean(point_iou_list)
    # point_avg_pa = np.mean(point_pa_list)

    # print(f"bbox_avg_iou: {bbox_avg_iou}")
    # print(f"bbox_avg_pa: {bbox_avg_pa}")
    # print(f"point_avg_iou: {point_avg_iou}")
    # print(f"point_avg_pa: {point_avg_pa}")
    
    # logging.info(f"bbox_avg_iou: {bbox_avg_iou}")
    # logging.info(f"bbox_avg_pa: {bbox_avg_pa}")
    # logging.info(f"point_avg_iou: {point_avg_iou}")
    # logging.info(f"point_avg_pa: {point_avg_pa}")
    
    
    bbox_avg_metrics = {}
    for metric_name, metric_list in bbox_metric_dicts.items():
        bbox_avg_metrics[metric_name] = np.mean(metric_list)
    for metric_name, avg_value in bbox_avg_metrics.items():
        print(f"bbox_{metric_name}: {avg_value}")
        logging.info(f"bbox_{metric_name}: {avg_value}")
    
    point_avg_metrics = {}
    for metric_name, metric_list in point_metric_dicts.items():
        point_avg_metrics[metric_name] = np.mean(metric_list)
    for metric_name, avg_value in point_avg_metrics.items():
        print(f"point_{metric_name}: {avg_value}")
        logging.info(f"point_{metric_name}: {avg_value}")
        
        
    log_performance_metrics(current_metrics, avg_metrics, logging)
    
if __name__ == "__main__":
    main()