自学教程：pytorch版本PSEnet训练并部署方式

torch==1.9.1torchvision==0.10.1python==3.8.0cuda==10.2mmcv==0.2.12editdistance==0.5.3Polygon3==3.0.9.1pyclipper==1.3.0opencv-python==3.4.2.17Cython==0.29.24

./compile.sh

import osfrom lxml import etreeimport numpy as npimport mathsrc_xml = "ANN"txt_dir = "gt"xml_listdir = os.listdir(src_xml)xml_listpath = [os.path.join(src_xml,xml_listdir1) for xml_listdir1 in xml_listdir]def xml_out(xml_path):    gt_lines = []    ET = etree.parse(xml_path)    objs = ET.findall("object")    for ix,obj in enumerate(objs):        name = obj.find("name").text        robox = obj.find("robndbox")        cx = int(float(robox.find("cx").text))        cy = int(float(robox.find("cy").text))        w = int(float(robox.find("w").text))        h = int(float(robox.find("h").text))        angle = float(robox.find("angle").text)        # angle = math.degrees(angle1)        wx1 = cx - int(0.5 * w)        wy1 = cy - int(0.5 * h)        wx2 = cx + int(0.5 * w)        wy2 = cy - int(0.5 * h)        wx3 = cx - int(0.5 * w)        wy3 = cy + int(0.5 * h)        wx4 = cx + int(0.5 * w)        wy4 = cy + int(0.5 * h)        x1 = int((wx1 - cx) * np.cos(angle) - (wy1 - cy) * np.sin(angle) + cx)        y1 = int((wx1 - cx) * np.sin(angle) - (wy1 - cy) * np.cos(angle) + cy)        x2 = int((wx2 - cx) * np.cos(angle) - (wy2 - cy) * np.sin(angle) + cx)        y2 = int((wx2 - cx) * np.sin(angle) - (wy2 - cy) * np.cos(angle) + cy)        x3 = int((wx3 - cx) * np.cos(angle) - (wy3 - cy) * np.sin(angle) + cx)        y3 = int((wx3 - cx) * np.sin(angle) - (wy3 - cy) * np.cos(angle) + cy)        x4 = int((wx4 - cx) * np.cos(angle) - (wy4 - cy) * np.sin(angle) + cx)        y4 = int((wx4 - cx) * np.sin(angle) - (wy4 - cy) * np.cos(angle) + cy)        lines = str(x1)+","+str(y1)+","+str(x2)+","+str(y2)+","+/                str(x3)+","+str(y3)+","+str(x4)+","+str(y4)+","+str(name)+"/n"        gt_lines.append(lines)        return gt_linesdef main():    count = 0    for xml_dir in xml_listdir:        gt_lines = xml_out(os.path.join(src_xml,xml_dir))        txt_path = "gt_" + xml_dir[:-4] + ".txt"        with open(os.path.join(txt_dir,txt_path),"a+") as fd:            fd.writelines(gt_lines)        count +=1        print("Write file %s" % str(count))if __name__ == "__main__":    main()

CUDA_VISIBLE_DEVICES=0,1,2,3 python train.py config/psenet/psenet_r50_ic15_736.py

import torchimport numpy as npimport argparseimport osimport os.path as ospimport sysimport timeimport jsonfrom mmcv import Configimport cv2from torchvision import transformsfrom dataset import build_data_loaderfrom models import build_modelfrom models.utils import fuse_modulefrom utils import ResultFormat, AverageMeterdef prepare_image(image, target_size):    """Do image preprocessing before prediction on any data.    :param image:       original image    :param target_size: target image size    :return:                        preprocessed image    """    #assert os.path.exists(img), 'file is not exists'    #img = cv2.imread(img)    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)    # h, w = image.shape[:2]    # scale = long_size / max(h, w)    img = cv2.resize(img, target_size)    # 将图片由(w,h)变为(1,img_channel,h,w)    tensor = transforms.ToTensor()(img)    tensor = tensor.unsqueeze_(0)    tensor = tensor.to(torch.device("cuda:0"))    return tensordef report_speed(outputs, speed_meters):    total_time = 0    for key in outputs:        if 'time' in key:            total_time += outputs[key]            speed_meters[key].update(outputs[key])            print('%s: %.4f' % (key, speed_meters[key].avg))    speed_meters['total_time'].update(total_time)    print('FPS: %.1f' % (1.0 / speed_meters['total_time'].avg))def load_model(cfg):    model = build_model(cfg.model)    model = model.cuda()    model.eval()    checkpoint = "psenet_r50_ic15_1024_finetune/checkpoint_580ep.pth.tar"    if checkpoint is not None:        if os.path.isfile(checkpoint):            print("Loading model and optimizer from checkpoint '{}'".format(checkpoint))            sys.stdout.flush()            checkpoint = torch.load(checkpoint)            d = dict()            for key, value in checkpoint['state_dict'].items():                tmp = key[7:]                d[tmp] = value            model.load_state_dict(d)        else:            print("No checkpoint found at")            raise        # fuse conv and bn    model = fuse_module(model)    return modelif __name__ == '__main__':    src_dir = "testimg/"    save_dir = "test_save/"    if not os.path.exists(save_dir):        os.makedirs(save_dir)    cfg = Config.fromfile("PSENet/config/psenet/psenet_r50_ic15_1024_finetune.py")    for d in [cfg, cfg.data.test]:        d.update(dict(            report_speed=False        ))    if cfg.report_speed:        speed_meters = dict(            backbone_time=AverageMeter(500),            neck_time=AverageMeter(500),            det_head_time=AverageMeter(500),            det_pse_time=AverageMeter(500),            rec_time=AverageMeter(500),            total_time=AverageMeter(500)        )    model = load_model(cfg)    model.eval()    count = 0    for img_name in os.listdir(src_dir):        img = cv2.imread(src_dir + img_name)        tensor = prepare_image(img, target_size=(1376, 1024))        data = dict()        img_metas = dict()        data['imgs'] = tensor        img_metas['org_img_size'] = torch.tensor([[img.shape[0], img.shape[1]]])        img_metas['img_size'] = torch.tensor([[1376, 1024]])        data['img_metas'] = img_metas        data.update(dict(            cfg=cfg        ))        with torch.no_grad():            outputs = model(**data)        if cfg.report_speed:            report_speed(outputs, speed_meters)        for bboxes in outputs['bboxes']:            x1 = bboxes[0]            y1 = bboxes[1]            x2 = bboxes[4]            y2 = bboxes[5]            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 3)        count = count + 1        cv2.imwrite(save_dir + img_name, img)        print("img test:", count)

from dataset import build_data_loaderfrom models import build_modelfrom models.utils import fuse_modulefrom utils import ResultFormat, AverageMeter