基于Unet+opencv实现天空对象的分割、替换和美化 - jsxyhelu

 getFiles("e:/template/Data_sky/data", fileNames);
    string saveFile = "e:/template/Data_sky/dataEX3/";
    for (int index = 0; index < fileNames.size(); index++)
    {
        Mat src = imread(fileNames[index]);
        Mat dst;
        string fileName;
        getFileName(fileNames[index], fileName);
        resize(src, dst, cv::Size(512, 512));
        imwrite(saveFile + fileName + "_512.jpg", dst);
        resize(src, dst, cv::Size(256, 256));
        imwrite(saveFile + fileName + "_256.jpg", dst);
        resize(src, dst, cv::Size(128, 128));
        imwrite(saveFile + fileName + "_128.jpg", dst);
        cout << fileName << endl;
    }
    fileNames.clear();
    getFiles("e:/template/Data_sky/mask", fileNames);
    saveFile = "e:/template/Data_sky/maskEX3/";
    for (int index = 0; index < fileNames.size(); index++)
    {
        Mat src = imread(fileNames[index], 0);
        Mat dst;
        string fileName;
        getFileName(fileNames[index], fileName);
        fileName = fileName.substr(0, fileName.size() - 3);
        resize(src, dst, cv::Size(512, 512));
        imwrite(saveFile + fileName + "_512_gt.jpg", dst);
        resize(src, dst, cv::Size(256, 256));
        imwrite(saveFile + fileName + "_256_gt.jpg", dst);
        resize(src, dst, cv::Size(128, 128));
        imwrite(saveFile + fileName + "_128_gt.jpg", dst);
        cout << fileName << endl;
    }

 # 5. Begin training
    for epoch in range(epochs):
        net.train()
        epoch_loss = 0
        with tqdm(total=n_train, desc=f'Epoch {epoch + 1}/{epochs}', unit='img') as pbar:
            for batch in train_loader:
                images = batch['image']
                true_masks = batch['mask']

                assert images.shape[1] == net.n_channels, \
                    f'Network has been defined with {net.n_channels} input channels, ' \
                    f'but loaded images have {images.shape[1]} channels. Please check that ' \
                    'the images are loaded correctly.'

                images = images.to(device=device, dtype=torch.float32)
                true_masks = true_masks.to(device=device, dtype=torch.long)
                ######
                one = torch.ones_like(true_masks)
                zero = torch.zeros_like(true_masks)
                true_masks = torch.where(true_masks>0,one,zero)
                #####
    
                with torch.cuda.amp.autocast(enabled=amp):
                    masks_pred = net(images)
                    loss = criterion(masks_pred, true_masks) \
                           + dice_loss(F.softmax(masks_pred, dim=1).float(),
                                       F.one_hot(true_masks, net.n_classes).permute(0, 3, 1, 2).float(),
                                       multiclass=True)

                optimizer.zero_grad(set_to_none=True)
                grad_scaler.scale(loss).backward()
                grad_scaler.step(optimizer)
                grad_scaler.update()

                pbar.update(images.shape[0])
                global_step += 1
                epoch_loss += loss.item()
                
                pbar.set_postfix(**{'loss (batch)': loss.item()})

                # Evaluation round
                division_step = (n_train // (10 * batch_size))
                if division_step > 0:
                    if global_step % division_step == 0:
                        histograms = {}
                        for tag, value in net.named_parameters():
                            tag = tag.replace('/', '.')
                           
                        val_score = evaluate(net, val_loader, device)
                        scheduler.step(val_score)

                        logging.info('Validation Dice score: {}'.format(val_score))

        if save_checkpoint:
            Path(dir_checkpoint).mkdir(parents=True, exist_ok=True)
            torch.save(net.state_dict(), str(dir_checkpoint / 'checkpoint_epoch{}.pth'.format(epoch + 1)))
            logging.info(f'Checkpoint {epoch + 1} saved!')

session = onnxruntime.InferenceSession("转换的onnx文件")
input_name = session.get_inputs()[0].name
label_name = session.get_outputs()[0].name

img_name_list = ['需要处理的图片']
image = Image.open(img_name_list[0])
w, h = image.size
dataset = SalObjDataset(
    img_name_list=img_name_list,
    lbl_name_list=[],
    transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)])
)
data_loader = DataLoader(
    dataset,
    batch_size=1,
    shuffle=False,
    num_workers=1
)
im = list(data_loader)[0]['image']
inputs_test = im
inputs_test = inputs_test.type(torch.FloatTensor)
with torch.no_grad():
    inputs_test = Variable(inputs_test)
res = session.run([label_name], {input_name: inputs_test.numpy().astype(np.float32)})
result = torch.from_numpy(res[0])
pred = result[:, 0, :, :]
pred = normPRED(pred)
pred = pred.squeeze()
predict_np = pred.cpu().data.numpy()
im = Image.fromarray(predict_np * 255).convert('RGB')
im = im.resize((w, h), resample=Image.BILINEAR)
im.show()

//    std::string strModelPath = "E:/template/u2net_train.pt";
void  bgr_u2net(cv::Mat& image_src, cv::Mat& result, torch::jit::Module& model)
{
    //1.模型已经导入
    auto device = torch::Device("cpu");
    //2.输入图片，变换到320
    cv::Mat  image_src1 = image_src.clone();
    cv::resize(image_src1, image_src1, cv::Size(320, 320));
    cv::cvtColor(image_src1, image_src1, cv::COLOR_BGR2RGB);
    // 3.图像转换为Tensor
    torch::Tensor tensor_image_src = torch::from_blob(image_src1.data, { image_src1.rows, image_src1.cols, 3 }, torch::kByte);
    tensor_image_src = tensor_image_src.permute({ 2,0,1 }); // RGB -> BGR互换
    tensor_image_src = tensor_image_src.toType(torch::kFloat);
    tensor_image_src = tensor_image_src.div(255);
    tensor_image_src = tensor_image_src.unsqueeze(0); // 拿掉第一个维度  [3, 320, 320]
    //4.网络前向计算
    auto src = tensor_image_src.to(device);
    auto pred = model.forward({ src }).toTuple()->elements()[0].toTensor();         //模型返回多个结果，用toTuple,其中elements()[i-1]获取第i个返回值                                                                                //d1,d2,d3,d4,d5,d6,d7= net(inputs_test) //pred = d1[:,0,:,:]
    auto res_tensor = (pred * torch::ones_like(src));
    res_tensor = normPRED(res_tensor);
    //是否就是Tensor转换为图像
    res_tensor = res_tensor.squeeze(0).detach();
    res_tensor = res_tensor.mul(255).clamp(0, 255).to(torch::kU8); //mul函数，表示张量中每个元素乘与一个数，clamp表示夹紧，限制在一个范围内输出
    res_tensor = res_tensor.to(torch::kCPU);
    //5.输出最终结果
    cv::Mat resultImg(res_tensor.size(1), res_tensor.size(2), CV_8UC3);
    std::memcpy((void*)resultImg.data, res_tensor.data_ptr(), sizeof(torch::kU8) * res_tensor.numel());
    cv::resize(resultImg, resultImg, cv::Size(image_src.cols, image_src.rows), cv::INTER_LINEAR);
    result = resultImg.clone();
}
 

int main()
{
    //参数和常量准备
    Net net = readNetFromONNX("E:/template/skyseg0113.onnx");
    if (net.empty()) {
        printf("read  model data failure...\n");
        return -1;
    }
    // load image data
    Mat frame = imread("E:\\sandbox/sky4.jpg");
    pyrDown(frame, frame);
    Mat blob;
    blobFromImage(frame, blob, 1.0 / 255.0, Size(320, 320), cv::Scalar(), true);
    net.setInput(blob);
    Mat prob = net.forward("473");
    Mat slice(cv::Size(prob.size[2], prob.size[3]), CV_32FC1, prob.ptr<float>(0, 0));
    threshold(slice, slice, 0.1, 1, cv::THRESH_BINARY_INV);
    normalize(slice, slice, 0, 255, NORM_MINMAX, CV_8U);
    
    Mat mask;
    resize(slice, mask, frame.size());//制作mask
    Mat matSrc = frame.clone();
    VP maxCountour = FindBigestContour(mask);
    Rect maxRect = boundingRect(maxCountour);
    if (maxRect.height == 0 || maxRect.width == 0)
        maxRect = Rect(0, 0, mask.cols, mask.rows);//特殊情况
    ////天空替换
    Mat matCloud = imread("E:/template/cloud/cloud1.jpg");
    resize(matCloud, matCloud, frame.size());
    //直接拷贝
    matCloud.copyTo(matSrc, mask);
    imshow("matSrc", matSrc);
    //seamless clone
    matSrc = frame.clone();
    Point center = Point((maxRect.x + maxRect.width) / 2, (maxRect.y + maxRect.height) / 2);//中间位置为蓝天的背景位置
    Mat normal_clone;
    Mat mixed_clone;
    Mat monochrome_clone;
    seamlessClone(matCloud, matSrc, mask, center, normal_clone, NORMAL_CLONE);
    seamlessClone(matCloud, matSrc, mask, center, mixed_clone, MIXED_CLONE);
    seamlessClone(matCloud, matSrc, mask, center, monochrome_clone, MONOCHROME_TRANSFER);
    imshow("normal_clone", normal_clone);
    imshow("mixed_clone", mixed_clone);
    imshow("monochrome_clone", monochrome_clone);
    waitKey();
    return 0;
}

int main()
{
    //参数和常量准备
    Net net = readNetFromONNX("E:/template/skyseg0113.onnx");
    if (net.empty()) {
        printf("read  model data failure...\n");
        return -1;
    }
    vector<string> vecFilePaths;
    getFiles("e:/template/sky", vecFilePaths);
    string strSavePath = "e:/template/sky_change_result";
    for (int index = 0;index<vecFilePaths.size();index++)
    {
        try{
            string strFilePath = vecFilePaths[index];
            string strFileName;
            getFileName(strFilePath, strFileName);
            Mat frame = imread(strFilePath);
            pyrDown(frame, frame);
            Mat blob;
            blobFromImage(frame, blob, 1.0 / 255.0, Size(320, 320), cv::Scalar(), true);
            net.setInput(blob);
            Mat prob = net.forward("473");
            Mat slice(cv::Size(prob.size[2], prob.size[3]), CV_32FC1, prob.ptr<float>(0, 0));
            threshold(slice, slice, 0.1, 1, cv::THRESH_BINARY_INV);
            normalize(slice, slice, 0, 255, NORM_MINMAX, CV_8U);
            Mat mask; 
            resize(slice, mask, frame.size());//制作mask
            Mat matSrc = frame.clone();
            VP maxCountour = FindBigestContour(mask);
            Rect maxRect = boundingRect(maxCountour);
            if (maxRect.height == 0 || maxRect.width == 0)
                maxRect = Rect(0, 0, mask.cols, mask.rows);//特殊情况
            Mat maskRedux(mask.size(), mask.type(), Scalar::all(0));
            Mat roi1 = mask(maxRect);
            Mat roi2 = maskRedux(maxRect);
            roi1.copyTo(roi2);
            ////天空替换
            Mat matCloud = imread("E:/template/cloud/cloud2.jpg");
            resize(matCloud, matCloud, frame.size());
            //直接拷贝
            matCloud.copyTo(matSrc, maskRedux);
            matSrc = frame.clone();
            cv::Point center = Point((maxRect.x + maxRect.width) / 2, (maxRect.y + maxRect.height) / 2);//中间位置为蓝天的背景位置
            Rect roi_s = maxRect;
            Rect roi_d(center.x - roi_s.width / 2, center.y - roi_s.height / 2, roi_s.width, roi_s.height);
            if(! (0 <= roi_d.x && 0 <= roi_d.width && roi_d.x + roi_d.width <= matSrc.cols && 0 <= roi_d.y && 0 <= roi_d.height && roi_d.y + roi_d.height <= matSrc.rows))
                center = Point(matSrc.cols / 2, matSrc.rows / 2);//这里错误的根源应该还是OpenCV 这块的代码有问题，其中roi_s不应该适用BoundingRect进行处理.所以采取补救的方法
            Mat mixed_clone;
            seamlessClone(matCloud, matSrc, maskRedux, center, mixed_clone, MIXED_CLONE);
            string saveFileName = strSavePath + "/" + strFileName + "_cloud2.jpg";
            imwrite(saveFileName, mixed_clone);
        }
        catch (Exception * e)
        {
            continue;
        }
    }

int main()
{
    Mat src = imread("e:/template/tiantan.jpg");
    Mat matCloud = imread("E:/template/cloud/cloud2.jpg");
    Mat mask = imread("e:/template/tiantanmask2.jpg", 0);
    resize(matCloud, matCloud, src.size());
    resize(mask, mask, src.size());
    Mat matSrc = src.clone();
    Mat board = mask.clone();
    cvtColor(board, board, COLOR_GRAY2BGR);
    //寻找模板最大轮廓
    VP maxCountour = FindBigestContour(mask);
    Rect maxRect = boundingRect(maxCountour);
    //异常处理
    Mat maskCopy = mask.clone();
    copyMakeBorder(maskCopy, maskCopy, 1, 1, 1, 1, BORDER_ISOLATED | BORDER_CONSTANT, Scalar(0));
    Rect roi_s = boundingRect(maskCopy);
    if (roi_s.empty()) return -1;
    cv::Point center = Point((maxRect.x + maxRect.width) / 2, (maxRect.y + maxRect.height) / 2);
    Rect roi_d(center.x - roi_s.width / 2, center.y - roi_s.height / 2, roi_s.width, roi_s.height);
    if (!(0 <= roi_d.x && 0 <= roi_d.width && roi_d.x + roi_d.width <= matSrc.cols && 0 <= roi_d.y && 0 <= roi_d.height && roi_d.y + roi_d.height <= matSrc.rows))
        center = Point(matSrc.cols / 2, matSrc.rows / 2);
    //融合
    Mat normal_clone, mixed_clone, monochrome_clone;
    seamlessClone(matCloud, matSrc, mask, center, normal_clone, NORMAL_CLONE);
    seamlessClone(matCloud, matSrc, mask, center, mixed_clone, MIXED_CLONE);
    seamlessClone(matCloud, matSrc, mask, center, monochrome_clone, MONOCHROME_TRANSFER);
    waitKey();
    return 0;
}