inference.py

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
from pycocotools.cocoeval import COCOeval
import json
import torch
from tqdm import tqdm
import numpy as np

"""
lyf: inference for 
"""
def inference_coco_RGB(dataset, model, model_type='Align_Dual_FCOS', threshold=0.05):
    """
    add img_rgb, img_ir 
    """
    model.eval()
    
    with torch.no_grad():

        # start collecting results
        results = []
        image_ids = []

        for index in range(len(dataset)):
            data = dataset[index]
            scale = data['scale']

            # run network
            if torch.cuda.is_available():   #[C, H, W] --> [1, C, H, W]
                scores, labels, boxes = model(data['img_rgb'].unsqueeze(dim=0).cuda().float(), data['boxes'].cuda(), data['classes'].cuda())
            else:
                scores, labels, boxes = model(data['img_rgb'].unsqueeze(dim=0).float(), data['boxes'], data['classes'])
                
            scores = scores.cuda()
            labels = labels.cuda()
            boxes  = boxes.cuda()

            # correct boxes for image scale
            boxes /= scale

            # correct boxes for image scale
            # change to (x, y, w, h) (MS COCO standard)
            boxes[:, :, 2] -= boxes[:, :, 0]
            boxes[:, :, 3] -= boxes[:, :, 1]


            # compute predicted labels and scores
            for box, score, label in zip(boxes[0], scores[0], labels[0]):
            #for box_id in range(boxes.shape[0]):
                #lyf: 
                label = int(label)
                # scores are sorted, so we can break
                if score < threshold:
                    break

                # append detection for each positively labeled class
                image_result = {
                        'image_id'    : generator.ids[index],
                        'category_id' : generator.id2category[label],  #lyf
                        'score'       : float(score),
                        'bbox'        : box.tolist(),
                }

                # append detection to results
                results.append(image_result)

            # append image to list of processed images
            image_ids.append(generator.ids[index])

            # print progress
            print('{}/{}'.format(index, len(generator)), end='\r')

        if not len(results):
            return

        # write output
        json.dump(results, open('{}_bbox_results.json'.format('VEDAI'), 'w'), indent=4)

        # load results in COCO evaluation tool
        coco_true = generator.coco
        coco_pred = coco_true.loadRes('{}_bbox_results.json'.format('VEDAI'))

        # run COCO evaluation
        coco_eval = COCOeval(coco_true, coco_pred, 'bbox')
        coco_eval.params.imgIds = image_ids
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()

        #coco_info = coco_eval.stats.tolist()  # numpy to list
        coco_stats, print_coco = summarize(coco_eval)


        if dataset.name == 'VEDAI':
            category_index = {0 : 'car', 1 : 'pickup', 2 : 'camping car', 3 : 'truck', 4 : 'other', 5 : 'tractor', 6 : 'boat', 7 : 'van'}
        elif dataset.name == 'FLIR':
            category_index = {0 : 'person', 1 : 'car', 2 : 'bicycle'}
        elif dataset.name == 'M3FD':
            category_index = {0: 'People', 1: 'Bus', 2: 'Car', 3: 'Motocycle', 4: 'Truck', 5: 'Lamp'}
        elif dataset.name == 'DroneVehicle':
            category_index = {0: 'car', 1: 'bus', 2: 'feright_car', 3: 'truck', 4: 'van'}
        # calculate voc info for every classes(IoU=0.5)
        voc_map_info_list = []
        for i in range(len(category_index)):
            stats, _ = summarize(coco_eval, catId=i)
            voc_map_info_list.append(" {:15}: {}".format(category_index[i], stats[1]))

        print_voc = "\n".join(voc_map_info_list)
        print(print_voc)
        
            # 将验证结果保存至txt文件中
        with open("./mAP_results/record_mAP_{}_{}.txt".format(dataset.name, model_type), "w") as f:
            record_lines = ["COCO results:",
                        print_coco,
                        "",
                        "mAP(IoU=0.5) for each category:",
                        print_voc]
            f.write("\n".join(record_lines))

        return 

def inference_coco_IR(dataset, model, model_type='Align_Dual_FCOS', threshold=0.05):
    """
    add img_rgb, img_ir 
    """
    model.eval()
    
    with torch.no_grad():

        # start collecting results
        results = []
        image_ids = []

        import time
        total_time = 0

        for index in range(len(dataset)):
            data = dataset[index]
            scale = data['scale']

            # run network
            if torch.cuda.is_available():   #[C, H, W] --> [1, C, H, W]
                start_time = time.time()
                scores, labels, boxes = model(data['img_ir'].unsqueeze(dim=0).cuda().float(), data['boxes'].cuda(), data['classes'].cuda())
                end_time = time.time()
            else:
                scores, labels, boxes = model(data['img_ir'].unsqueeze(dim=0).float(), data['boxes'], data['classes'])


            total_time += end_time - start_time  

            scores = scores.cuda()
            labels = labels.cuda()
            boxes  = boxes.cuda()

            # correct boxes for image scale
            boxes /= scale

            # correct boxes for image scale
            # change to (x, y, w, h) (MS COCO standard)
            boxes[:, :, 2] -= boxes[:, :, 0]
            boxes[:, :, 3] -= boxes[:, :, 1]

            

            # compute predicted labels and scores
            for box, score, label in zip(boxes[0], scores[0], labels[0]):
            #for box_id in range(boxes.shape[0]):
                #lyf: 
                label = int(label)
                # scores are sorted, so we can break
                if score < threshold:
                    break

                # append detection for each positively labeled class
                image_result = {
                        'image_id'    : generator.ids[index],
                        'category_id' : generator.id2category[label],  #lyf
                        'score'       : float(score),
                        'bbox'        : box.tolist(),
                }

                # append detection to results
                results.append(image_result)

            # append image to list of processed images
            image_ids.append(generator.ids[index])

            # print progress
            print('{}/{}'.format(index, len(generator)), end='\r')

        average_time = total_time / len(dataset)
        print(f"Average inference time per iteration: {average_time} seconds")

        if not len(results):
            return

        # write output
        json.dump(results, open('{}_bbox_results.json'.format('VEDAI'), 'w'), indent=4)

        # load results in COCO evaluation tool
        coco_true = generator.coco
        coco_pred = coco_true.loadRes('{}_bbox_results.json'.format('VEDAI'))

        # run COCO evaluation
        coco_eval = COCOeval(coco_true, coco_pred, 'bbox')
        coco_eval.params.imgIds = image_ids
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()

        #coco_info = coco_eval.stats.tolist()  # numpy to list
        coco_stats, print_coco = summarize(coco_eval)

        if dataset.name == 'VEDAI':
            category_index = {0 : 'car', 1 : 'pickup', 2 : 'camping car', 3 : 'truck', 4 : 'other', 5 : 'tractor', 6 : 'boat', 7 : 'van'}
        elif dataset.name == 'FLIR':
            category_index = {0 : 'person', 1 : 'car', 2 : 'bicycle'}
        elif dataset.name == 'M3FD':
            category_index = {0: 'People', 1: 'Bus', 2: 'Car', 3: 'Motocycle', 4: 'Truck', 5: 'Lamp'}
        elif dataset.name == 'DroneVehicle':
            category_index = {0: 'car', 1: 'bus', 2: 'feright_car', 3: 'truck', 4: 'van'}

        # calculate voc info for every classes(IoU=0.5)
        voc_map_info_list = []
        for i in range(len(category_index)):
            stats, _ = summarize(coco_eval, catId=i)
            voc_map_info_list.append(" {:15}: {}".format(category_index[i], stats[1]))

        print_voc = "\n".join(voc_map_info_list)
        print(print_voc)
        
            # 将验证结果保存至txt文件中
        with open("./mAP_results/record_mAP_{}_{}.txt".format(dataset.name, model_type), "w") as f:
            record_lines = ["COCO results:",
                        print_coco,
                        "",
                        "mAP(IoU=0.5) for each category:",
                        print_voc]
            f.write("\n".join(record_lines))

        return 
    

def inference_coco_multispectral(dataset, model, model_type='Align_Dual_FCOS', threshold=0.05):
    """
    add img_rgb, img_ir 
    """
    model.eval()
    
    with torch.no_grad():

        # start collecting results
        results = []
        image_ids = []


        import time
        total_time = 0
        
        for index in range(len(dataset)):
            data = dataset[index]
            scale = data['scale']
            
            # run network
            if torch.cuda.is_available():   #[C, H, W] --> [1, C, H, W]
                
                start_time = time.time()
                scores, labels, boxes = model(data['img_rgb'].unsqueeze(dim=0).cuda().float(), data['img_ir'].unsqueeze(dim=0).cuda().float(), data['boxes'].cuda(), data['classes'].cuda())
                end_time = time.time()  
            else:
                scores, labels, boxes = model(data['img_rgb'].unsqueeze(dim=0).float(), data['img_ir'].unsqueeze(dim=0).float(), data['boxes'], data['classes'])
            
            total_time += end_time - start_time  

            scores = scores.cuda()
            labels = labels.cuda()
            boxes  = boxes.cuda()

            # correct boxes for image scale
            boxes /= scale

            # correct boxes for image scale
            # change to (x, y, w, h) (MS COCO standard)
            boxes[:, :, 2] -= boxes[:, :, 0]
            boxes[:, :, 3] -= boxes[:, :, 1]


            # compute predicted labels and scores
            for box, score, label in zip(boxes[0], scores[0], labels[0]):
            #for box_id in range(boxes.shape[0]):
                #lyf: 
                label = int(label)
                # scores are sorted, so we can break
                if score < threshold:
                    break

                # append detection for each positively labeled class
                image_result = {
                        'image_id'    : generator.ids[index],
                        'category_id' : generator.id2category[label],  #lyf
                        'score'       : float(score),
                        'bbox'        : box.tolist(),
                }

                # append detection to results
                results.append(image_result)

            # append image to list of processed images
            image_ids.append(generator.ids[index])

            # print progress
            print('{}/{}'.format(index, len(generator)), end='\r')

        average_time = total_time / len(dataset)
        print(f"Average inference time per iteration: {average_time} seconds")
    
        if not len(results):
            return

        # write output
        json.dump(results, open('{}_bbox_results.json'.format('VEDAI'), 'w'), indent=4)

        # load results in COCO evaluation tool
        coco_true = generator.coco
        coco_pred = coco_true.loadRes('{}_bbox_results.json'.format('VEDAI'))

        # run COCO evaluation
        coco_eval = COCOeval(coco_true, coco_pred, 'bbox')
        coco_eval.params.imgIds = image_ids
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()

        #coco_info = coco_eval.stats.tolist()  # numpy to list
        coco_stats, print_coco = summarize(coco_eval)

        if dataset.name == 'VEDAI':
            category_index = {0 : 'car', 1 : 'pickup', 2 : 'camping car', 3 : 'truck', 4 : 'other', 5 : 'tractor', 6 : 'boat', 7 : 'van'}
        elif dataset.name == 'FLIR':
            category_index = {0 : 'person', 1 : 'car', 2 : 'bicycle'}
        elif dataset.name == 'M3FD':
            category_index = {0: 'People', 1: 'Bus', 2: 'Car', 3: 'Motocycle', 4: 'Truck', 5: 'Lamp'}
        elif dataset.name == 'DroneVehicle':
            category_index = {0: 'car', 1: 'bus', 2: 'feright_car', 3: 'truck', 4: 'van'}

        # calculate voc info for every classes(IoU=0.5)
        voc_map_info_list = []
        for i in range(len(category_index)):
            stats, _ = summarize(coco_eval, catId=i)
            voc_map_info_list.append(" {:15}: {}".format(category_index[i], stats[1]))

        print_voc = "\n".join(voc_map_info_list)
        print(print_voc)
        
            # 将验证结果保存至txt文件中
        with open("./mAP_results/record_mAP_{}_{}.txt".format(dataset.name, model_type), "w") as f:
            record_lines = ["COCO results:",
                        print_coco,
                        "",
                        "mAP(IoU=0.5) for each category:",
                        print_voc]
            f.write("\n".join(record_lines))

        return 
    

def summarize(self, catId=None):
    """
    Compute and display summary metrics for evaluation results.
    Note this functin can *only* be applied on the default parameter setting
    """

    def _summarize(ap=1, iouThr=None, areaRng='all', maxDets=100):
        p = self.params
        iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}'
        titleStr = 'Average Precision' if ap == 1 else 'Average Recall'
        typeStr = '(AP)' if ap == 1 else '(AR)'
        iouStr = '{:0.2f}:{:0.2f}'.format(p.iouThrs[0], p.iouThrs[-1]) \
            if iouThr is None else '{:0.2f}'.format(iouThr)

        aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
        mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]

        if ap == 1:
            # dimension of precision: [TxRxKxAxM]
            s = self.eval['precision']
            # IoU
            if iouThr is not None:
                t = np.where(iouThr == p.iouThrs)[0]
                s = s[t]

            if isinstance(catId, int):
                s = s[:, :, catId, aind, mind]
            else:
                s = s[:, :, :, aind, mind]

        else:
            # dimension of recall: [TxKxAxM]
            s = self.eval['recall']
            if iouThr is not None:
                t = np.where(iouThr == p.iouThrs)[0]
                s = s[t]

            if isinstance(catId, int):
                s = s[:, catId, aind, mind]
            else:
                s = s[:, :, aind, mind]

        if len(s[s > -1]) == 0:
            mean_s = -1
        else:
            mean_s = np.mean(s[s > -1])

        print_string = iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s)
        return mean_s, print_string

    stats, print_list = [0] * 12, [""] * 12
    stats[0], print_list[0] = _summarize(1)
    stats[1], print_list[1] = _summarize(1, iouThr=.5, maxDets=self.params.maxDets[2])
    stats[2], print_list[2] = _summarize(1, iouThr=.75, maxDets=self.params.maxDets[2])
    stats[3], print_list[3] = _summarize(1, areaRng='small', maxDets=self.params.maxDets[2])
    stats[4], print_list[4] = _summarize(1, areaRng='medium', maxDets=self.params.maxDets[2])
    stats[5], print_list[5] = _summarize(1, areaRng='large', maxDets=self.params.maxDets[2])
    stats[6], print_list[6] = _summarize(0, maxDets=self.params.maxDets[0])
    stats[7], print_list[7] = _summarize(0, maxDets=self.params.maxDets[1])
    stats[8], print_list[8] = _summarize(0, maxDets=self.params.maxDets[2])
    stats[9], print_list[9] = _summarize(0, areaRng='small', maxDets=self.params.maxDets[2])
    stats[10], print_list[10] = _summarize(0, areaRng='medium', maxDets=self.params.maxDets[2])
    stats[11], print_list[11] = _summarize(0, areaRng='large', maxDets=self.params.maxDets[2])

    print_info = "\n".join(print_list)

    if not self.eval:
        raise Exception('Please run accumulate() first')

    return stats, print_info

if __name__ == "__main__":
    from dataset.VEDAI_dataset import VEDAIGenerator
    # generator=VEDAIGenerator("/data1/users/liuyanfeng/MultispectralOD/VEDAI_COCO/test/",
    #                         "/data1/users/liuyanfeng/MultispectralOD/VEDAI_COCO/annotations/test_coco_format.json",)
    #from model.fcos import FCOSDetector
    #model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    #model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/VEDAI/baseline_IR/VEDAI_FCOS_0.7533544149385896.pth"))
    #inference_coco_RGB(generator, model, model_type="FCOS_RGB")

    #model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    #model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/VEDAI/baseline_RGB/VEDAI_FCOS_0.7773923004053019.pth"))
    #inference_coco_RGB(generator, model, model_type="FCOS_RGB")

    # from model.dual_fcos import Dual_FCOSDetector
    # model = Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/VEDAI/baseline_RGB+IR/VEDAI_FCOS_0.7818216981522094.pth"))
    # inference_coco_multispectral(generator, model, model_type='Dual_FCOS')

    # from model.Align_Dual_FCOS import Align_Dual_FCOSDetector
    # model = Align_Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/VEDAI/baseline+align_fpn+align_head+dual_align/VEDAI_FCOS_0.8091348336519473.pth"))
    # inference_coco_multispectral(generator, model, model_type="Align_Dual_FCOS")
    
    #from model.fcos_PVT2 import FCOSDetector
    #model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    #model.load_state_dict(torch.load("./save_weights/VEDAI/baseline_IR_PVT2/VEDAI_FCOS_0.7571636531540242.pth"))
    #inference_coco_IR(generator, model, model_type="FCOS_IR_PVT2")
    
    #model.load_state_dict(torch.load("./save_weights/VEDAI/baseline_RGB_PVT2/VEDAI_FCOS_0.7913594968837172.pth"))
    #inference_coco_RGB(generator, model, model_type="FCOS_RGB_PVT2")

    # from model.dual_fcos_PVT2 import Dual_FCOSDetector
    # model = Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/VEDAI/baseline_RGB+IR_PVT2/VEDAI_FCOS_0.7987801996849406.pth"))
    # inference_coco_multispectral(generator, model, model_type='Dual_FCOS_PVT2')

    #from model.Align_Dual_FCOS_PVT2 import Align_Dual_FCOSDetector
    #model = Align_Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    #model.load_state_dict(torch.load("./save_weights/VEDAI/baseline+align_fpn+align_head+dual_align_PVT/VEDAI_FCOS_0.8343130425154863.pth"))
    #inference_coco_multispectral(generator, model, model_type='Align_Dual_FCOS_PVT2')


    from dataset.DroneVehicle_dataset import DroneVehicleGenerator
    # from model.fcos import FCOSDetector
    generator = DroneVehicleGenerator()
    # model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/DroneVehicle/baseline_RGB/DroneVehicle_FCOS_0.6113890697301754.pth"))
    # inference_coco_RGB(generator, model, model_type="FCOS_RGB")
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/DroneVehicle/baseline_IR/DroneVehicle_FCOS_0.6573492082687165.pth"))
    # inference_coco_IR(generator, model, model_type="FCOS_IR")

    from model.fcos_PVT import FCOSDetector
    model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    #model.load_state_dict(torch.load("./save_weights/DroneVehicle/baseline_RGB_PVT/DroneVehicle_FCOS_0.6471979759214866.pth"))
    #inference_coco_RGB(generator, model, model_type="FCOS_RGB_PVT")
    #model.load_state_dict(torch.load("./save_weights/DroneVehicle/baseline_IR_PVT/DroneVehicle_FCOS_0.7210628429796552.pth"))
    #inference_coco_IR(generator, model, model_type="FCOS_IR_PVT")

    # from model.dual_fcos import Dual_FCOSDetector
    # model = Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/DroneVehicle/baseline_RGB+IR/DroneVehicle_FCOS_0.70556.pth"))
    # inference_coco_multispectral(generator, model, model_type="Dual_FCOS")

    # from model.Align_Dual_FCOS_residual import Align_Dual_FCOSDetector
    # model = Align_Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/DroneVehicle/baseline+align_head+dual_align/DroneVehicle_FCOS_0.7255275681279547.pth"))
    # inference_coco_multispectral(generator, model, model_type="Align_Dual_FCOS")
    
    # from dataset.DroneVehicle_dataset import DroneVehicleGenerator
    # generator = DroneVehicleGenerator()
    # from model.Align_Dual_FCOS_residual_PVT import Align_Dual_FCOSDetector
    # model = Align_Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/DroneVehicle/baseline+align_head+dual_align_PVT/DroneVehicle_FCOS_0.7490874383633909.pth"))
    # inference_coco_multispectral(generator, model, model_type="Align_Dual_FCOS_PVT")


    from dataset.FLIR_Aligned_dataset import FLIRGenerator
    # generator = FLIRGenerator()

    # from model.Align_Dual_FCOS import Align_Dual_FCOSDetector
    # model = Align_Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/FLIR/baseline_dual_align_head_align/FLIR_FCOS_0.7463164857249269.pth"))#"/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/FLIR/baseline_dual_align_head_align/FLIR_FCOS_0.7435085614628874.pth"))
    # inference_coco_multispectral(generator,model, model_type="Align_Dual_FCOS")

    #from model.fcos_PVT import FCOSDetector
    # model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/FLIR/baseline_RGB_PVT/FLIR_FCOS_0.6007557062171366.pth"))
    # inference_coco_RGB(generator,model, model_type="FCOS_RGB_PVT")

    # model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/FLIR/baseline_IR_PVT/FLIR_FCOS_0.7048883905333573.pth"))
    # inference_coco_IR(generator,model, model_type="FCOS_IR_PVT")
    
    # from model.dual_fcos_PVT import Dual_FCOSDetector
    # model = Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/FLIR/baseline_RGB+IR_PVT/FLIR_FCOS_0.7342263407087976.pth"))
    # inference_coco_multispectral(generator,model, model_type="Dual_FCOS_PVT")
    # from model.Align_Dual_FCOS_PVT import Align_Dual_FCOSDetector
    # generator = FLIRGenerator()
    # model = Align_Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/FLIR/baseline_dual_align_head_align_PVT/FLIR_FCOS_0.7594769071599681.pth"))
    # inference_coco_multispectral(generator,model, model_type="Align_Dual_FCOS_PVT")


    # from model.dual_fcos import Dual_FCOSDetector
    # model = Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/FLIR/baseline_RGB+IR/FLIR_FCOS_0.7224278987521715.pth"))
    # inference_coco_multispectral(generator,model, model_type="Dual_FCOS")

    # # from model.fcos import FCOSDetector
    # # model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/FLIR/baseline_RGB/FLIR_FCOS_0.6060130727113124.pth"))
    # # inference_coco_RGB(generator,model, model_type="FCOS_RGB")

    # from model.fcos import FCOSDetector
    # model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()    
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/FLIR/baseline_IR/FLIR_FCOS_0.6687149733204729.pth"))
    # inference_coco_IR(generator,model, model_type="FCOS_IR")

    from dataset.M3FD_dataset import M3FDGenerator
    # generator = M3FDGenerator()

    # from model.fcos import FCOSDetector
    # model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/M3FD/baseline_RGB/M3FD_FCOS_0.525124566344887.pth"))
    # inference_coco_RGB(generator, model, model_type="FCOS_RGB")
    # from thop import profile
    # input = torch.randn(1, 3, 640, 512).cuda()

    # from thop import clever_format
    # macs, params = profile(model, inputs=(input,None, None))
    # macs, params = clever_format((macs, params), "%.3f")
    # print(macs)  # 打印格式化后的浮点运算次数
    # print(params)
    # from model.fcos import FCOSDetector
    # model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/M3FD/baseline_IR/M3FD_FCOS_0.5048685918131055.pth"))
    # inference_coco_IR(generator, model, model_type="FCOS_IR")

    # from model.dual_fcos import Dual_FCOSDetector
    # model = Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/M3FD/baseline_RGB+IR/M3FD_FCOS_0.5779541788702803.pth"))
    # inference_coco_multispectral(generator, model, model_type="Dual_FCOS")
    
    # from thop import profile

    # from thop import clever_format
    # macs, params = profile(model, inputs=(input,input,None, None))
    # macs, params = clever_format((macs, params), "%.3f")
    # print(macs)  # 打印格式化后的浮点运算次数
    # print(params)
    # inference_coco_multispectral(generator, model, model_type="Dual_FCOS")

    

    # from thop import clever_format
    # macs, params = profile(model, inputs=(input,input,None, None))
    # macs, params = clever_format((macs, params), "%.3f")
    # print(macs)  # 打印格式化后的浮点运算次数
    # print(params)


    #from model.fcos_PVT import FCOSDetector
    # model = FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/M3FD/baseline_IR_PVT/M3FD_FCOS_0.4903807750261591.pth"))
    # inference_coco_IR(generator, model, model_type="FCOS_IR_PVT")
    # model.load_state_dict(torch.load("./save_weights/M3FD/baseline_RGB_PVT/M3FD_FCOS_0.5352943617329038.pth"))
    # inference_coco_RGB(generator, model, model_type="FCOS_RGB_PVT")
    # from model.dual_fcos_PVT import Dual_FCOSDetector
    # model = Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/M3FD/baseline_RGB+IR_PVT/M3FD_FCOS_0.5832574280826551.pth"))
    # inference_coco_multispectral(generator,model, model_type="Dual_FCOS_PVT")

    # from model.Align_Dual_FCOS_residual import Align_Dual_FCOSDetector
    # model = Align_Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("/data1/users/liuyanfeng/MultispectralOD/FCOS/save_weights/M3FD/baseline_dual_align_head_align/M3FD_FCOS_0.5947070011745564.pth"))
    # inference_coco_multispectral(generator, model, model_type="Align_Dual_FCOS")
   
    # from model.Align_Dual_FCOS_residual_PVT import Align_Dual_FCOSDetector
    # model = Align_Dual_FCOSDetector(num_class=len(generator.category2id.keys())).cuda().eval()
    # model.load_state_dict(torch.load("./save_weights/M3FD/baseline_dual_align_head_align_PVT/M3FD_FCOS_0.6173570051152696.pth"))
    # inference_coco_multispectral(generator, model, model_type="Align_Dual_FCOS_PVT")



    from dataset.VEDAI_dataset import VEDAIGenerator
    generator = VEDAIGenerator()

    


"""
FCOS: 
64.532G
32.123M

Dual_FCOS:
164.010G
154.732M

Align_Dual_FCOS:
149.555G
96.754M
"""



"""

Average inference time per iteration: 0.13926207842737406 seconds  选择 for dual-modal FCOS 

Average inference time per iteration: 0.12068611848271964 seconds 选择 for single-modal FCOS 

"""


"""
Average inference time per iteration: 0.16874815822236042 seconds  选择 for Align_Dual_FCOS 
"""





"""
Average inference time per iteration: 0.0922149556944269 seconds 选择 for single-modal FCOS PVT 
"""


"""
Average inference time per iteration: 0.0971854863479791 seconds 选择 for dual-modal FCOS PVT 
"""


"""
Average inference time per iteration: 0.12616108596094786 seconds 选择 for Align_Dual_FCOS_PVT
"""