【复现代码】调试 demo.py网站首页 学无止境

【复现代码】调试 demo.py

dearRongerr 2025-04-11 00:01:02

简介【复现代码】调试 demo.py

from torch.nn import DataParallel
from models.dave import build_model
from utils.arg_parser import get_argparser
import argparse
import torch
import os
import matplotlib.patches as patches
from PIL import Image

from utils.data import resize
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings("ignore")


bounding_boxes = []


def on_click(event):
    # Record the starting point of the bounding box
    global ix, iy
    ix, iy = event.xdata, event.ydata
    # Connect the release event
    fig.canvas.mpl_connect('button_release_event', on_release)


def on_release(event):
    # Record the ending point of the bounding box
    global ix, iy
    x, y = event.xdata, event.ydata
    # Calculate the width and height of the bounding box
    width = x - ix
    height = y - iy
    # Add a rectangle patch to the axes
    rect = patches.Rectangle((ix, iy), width, height, edgecolor='r', facecolor='none')
    ax.add_patch(rect)
    # Store the bounding box coordinates
    bounding_boxes.append((ix, iy, ix + width, iy + height))
    plt.draw()


@torch.no_grad()
def demo(args):
    img_path = "material//458.jpg"
    global fig, ax

    gpu = 0
    # torch.cuda.set_device(gpu) 改
    # device = torch.device('cpu') # device = torch.device(gpu) 改

    model = DataParallel(
        build_model(args)   # build_model(args).to(device),改
        # device_ids=[gpu],
        # output_device=gpu
    )
    model.load_state_dict(
        torch.load(os.path.join(args.model_path, 'DAVE_3_shot.pth'),
                   map_location=torch.device('cpu'),
                   weights_only=True)['model'], 
        strict=False       
    )
    pretrained_dict_feat = {k.split("feat_comp.")[1]: v for k, v in
                            torch.load(os.path.join(args.model_path, 'verification.pth'),
                                       map_location=torch.device('cpu')
                                       )['model'].items() if 'feat_comp' in k}
    model.module.feat_comp.load_state_dict(pretrained_dict_feat)
    model.eval()

    image = Image.open(img_path).convert("RGB")
    # Create a figure and axis
    fig, ax = plt.subplots(1)
    ax.imshow(image)
    # Connect the click event
    cid = fig.canvas.mpl_connect('button_press_event', on_click)
    plt.title("Click and drag to draw bboxes, then close window")
    # Show the image
    plt.show()

    bboxes = torch.tensor(bounding_boxes)

    img, bboxes, scale = resize(image, bboxes)
    img = img.unsqueeze(0) # 改 img = img.unsqueeze(0).to(device)
    bboxes = bboxes.unsqueeze(0) # bboxes = bboxes.unsqueeze(0).to(device)

    denisty_map, _, tblr, predicted_bboxes = model(img, bboxes=bboxes)

    plt.clf()
    plt.imshow(image)
    pred_boxes = predicted_bboxes.box.cpu() / torch.tensor([scale[0], scale[1], scale[0], scale[1]])
    for i in range(len(pred_boxes)):
        box = pred_boxes[i]

        plt.plot([box[0], box[0], box[2], box[2], box[0]], [box[1], box[3], box[3], box[1], box[1]], linewidth=2,
                 color='red')
    plt.title("Dmap count:" + str(round(denisty_map.sum().item(), 1)) + " Box count:" + str(len(pred_boxes)))
    plt.show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser('DAVE', parents=[get_argparser()])
    args = parser.parse_args()
    demo(args)

运行是没有问题的，运行命令：

python demo.py --skip_train --model_name DAVE_3_shot --model_path material --backbone resnet50 --swav_backbone --reduction 8 --num_enc_layers 3 --num_dec_layers 3 --kernel_dim 3 --emb_dim 256 --num_objects 3 --num_workers 8 --use_query_pos_emb --use_objectness --use_appearance --batch_size 1 --pre_norm

python demo.py

--skip_train

--model_name DAVE_3_shot

--model_path material

--backbone resnet50

--swav_backbone

--reduction 8

--num_enc_layers 3

--num_dec_layers 3

--kernel_dim 3

--emb_dim 256

--num_objects 3

--num_workers 8

--use_query_pos_emb

--use_objectness

--use_appearance

--batch_size 1

--pre_norm

问题的是

import pdb;pdb.set_trace()

依然启动不了调试，（我真不会，崩溃）

之前 main.py的参数设置：

Namespace(

aux_weight=0.3,

backbone='resnet18', backbone_lr=0, batch_size=4, count_loss_weight=0, d_s=1.0, d_t=3, data_path='data', dataset='fsc147', det_model_name='DAVE', det_train=False, detection_loss_weight=0.01, dropout=0.1, egv=0.132, emb_dim=256, epochs=200, eval_multicat=False, fcos_pred_size=512, i_thr=0.55, image_size=512, kernel_dim=3, lr=0.0001, lr_drop=200, m_s=0.0, max_grad_norm=0.1, min_count_loss_weight=0, model_name='DAVE_3_shot', model_path='material/', norm_s=False, normalized_l2=False, num_dec_layers=3, num_enc_layers=3, num_heads=8, num_objects=3, num_workers=0, orig_dmaps=False, pre_norm=False, prompt_shot=False, reduction=8, resume_training=False, s_t=0.008, skip_cars=False, skip_test=False, skip_train=False, swav_backbone=False, task='fscd147', tiling_p=0.5, unseen=False, use_appearance=False, use_objectness=False, use_query_pos_emb=False, weight_decay=0.0001, zero_shot=False)

算了，看看 demo 里哪里需要设置参数了，为什么直接运行会出错。