app.py

# -*- coding: utf-8 -*-

import sys
import io
import requests
import json
import base64
from PIL import Image
import numpy as np
import gradio as gr
import mmengine
from mmengine import Config, get

import argparse
import os
import cv2
import yaml
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import datasets
import models
import numpy as np

from torchvision import transforms
from mmcv.runner import load_checkpoint
import visual_utils
from PIL import Image
from models.utils_prompt import get_prompt_inp, pre_prompt, pre_scatter_prompt, get_prompt_inp_scatter

device = torch.device("cpu")

def batched_predict(model, inp, coord, bsize):
    with torch.no_grad():
        model.gen_feat(inp)
        n = coord.shape[1]
        ql = 0
        preds = []
        while ql < n:
            qr = min(ql + bsize, n)
            pred = model.query_rgb(coord[:, ql: qr, :])
            preds.append(pred)
            ql = qr
        pred = torch.cat(preds, dim=1)
    return pred, preds


def tensor2PIL(tensor):
    toPIL = transforms.ToPILImage()
    return toPIL(tensor)


def Decoder1_optical_instance(image_input):
    with open('configs/fine_tuning_one_decoder.yaml', 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    model = models.make(config['model']).cpu()
    sam_checkpoint = torch.load("./save/model_epoch_last.pth", map_location='cpu')
    model.load_state_dict(sam_checkpoint, strict=False)
    model.eval()

    # img = np.array(image_input).copy()
    label2color = visual_utils.Label2Color(cmap=visual_utils.color_map('Unify_double'))
    # image_input.save(f'./save/visual_fair1m/input_img.png', quality=5)
    img = transforms.Resize([1024, 1024])(image_input)
    transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229,0.224,0.225])])
    input_img = transform(img)
    input_img = input_img.unsqueeze(0)
    image_embedding = model.image_encoder(input_img) # torch.Size([1, 256, 64, 64])
    sparse_embeddings, dense_embeddings, scatter_embeddings  = model.prompt_encoder(
        points=None,
        boxes=None,
        masks=None,
    scatter=None)
    # 目标类预测decoder
    low_res_masks, iou_predictions = model.mask_decoder(
        image_embeddings=image_embedding,
        image_pe=model.prompt_encoder.get_dense_pe(),
        sparse_prompt_embeddings=sparse_embeddings,
        dense_prompt_embeddings=dense_embeddings,
        multimask_output=False
    )
    pred = model.postprocess_masks(low_res_masks, model.inp_size, model.inp_size)
    _, prediction = pred.max(dim=1)
    prediction_to_save = label2color(prediction.cpu().numpy().astype(np.uint8))[0]

    return prediction_to_save


def Decoder1_optical_terrain(image_input):
    with open('configs/fine_tuning_one_decoder.yaml', 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    model = models.make(config['model']).cpu()
    sam_checkpoint = torch.load("./save/model_epoch_last.pth", map_location='cpu')
    model.load_state_dict(sam_checkpoint, strict=False)
    model.eval()

    denorm = visual_utils.Denormalize(mean=[0.485, 0.456, 0.406],std=[0.229,0.224,0.225])
    label2color = visual_utils.Label2Color(cmap=visual_utils.color_map('Unify_Vai'))
    # image_input.save(f'./save/visual_fair1m/input_img.png', quality=5)
    img = transforms.Resize([1024, 1024])(image_input)
    transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229,0.224,0.225])])
    input_img = transform(img)
    input_img = torch.unsqueeze(input_img, dim=0)
    # input_img = transforms.ToTensor()(img).unsqueeze(0)
    image_embedding = model.image_encoder(input_img) # torch.Size([1, 256, 64, 64])
    sparse_embeddings, dense_embeddings, scatter_embeddings = model.prompt_encoder(
        points=None,
        boxes=None,
        masks=None,
        scatter=None)
    low_res_masks_instanse, iou_predictions = model.mask_decoder(
        image_embeddings=image_embedding,
        # image_embeddings=image_embedding.unsqueeze(0),
        image_pe=model.prompt_encoder.get_dense_pe(),
        sparse_prompt_embeddings=sparse_embeddings,
        dense_prompt_embeddings=dense_embeddings,
        # multimask_output=multimask_output,
        multimask_output=False
    )
    # 地物类预测decoder
    low_res_masks, iou_predictions_2 = model.mask_decoder_diwu(
        image_embeddings=image_embedding,
        image_pe=model.prompt_encoder.get_dense_pe(),
        sparse_prompt_embeddings=sparse_embeddings,
        dense_prompt_embeddings=dense_embeddings,
        # multimask_output=False,
        multimask_output=True,
    )  # B*C+1*H*W

    pred_instance = model.postprocess_masks(low_res_masks_instanse, model.inp_size, model.inp_size)
    pred = model.postprocess_masks(low_res_masks, model.inp_size, model.inp_size)
    pred = torch.softmax(pred,dim=1)
    pred_instance = torch.softmax(pred_instance,dim=1)
    _, prediction = pred.max(dim=1)
    prediction[prediction==12]=0  #把第二个decoder里得背景变成0
    print(torch.unique(prediction))
    _, prediction_instance = pred_instance.max(dim=1)
    print(torch.unique(prediction_instance))
    prediction_sum = prediction + prediction_instance #没有冲突的位置就会正常猜测
    print(torch.unique(prediction_sum))
    prediction_tmp = prediction_sum.clone()
    prediction_tmp[prediction_tmp==1] = 255
    prediction_tmp[prediction_tmp==2] = 255
    prediction_tmp[prediction_tmp==5] = 255
    prediction_tmp[prediction_tmp==6] = 255
    prediction_tmp[prediction_tmp==14] = 255
    # prediction_tmp[prediction_tmp==0] = 255 #同时是背景
    # index = prediction_tmp != 255
    pred[:, 0][prediction_tmp == 255]=100 #把已经决定的像素位置的背景预测概率设置为最大
    pred_instance[:, 0][prediction_tmp == 255]=100#把已经决定的像素位置的背景预测概率设置为最大
    buchong = torch.zeros([1,2,1024,1024])
    pred = torch.cat((pred, buchong),dim=1)
    # print(torch.unique(torch.argmax(pred,dim=1)))
    # Decoder1_logits = torch.zeros([1,15,1024,1024]).cuda()
    Decoder2_logits = torch.zeros([1,15,1024,1024])
    Decoder2_logits[:,0,...] = pred[:,0,...]
    Decoder2_logits[:,5,...] = pred_instance[:,5,...]
    Decoder2_logits[:,14,...] = pred_instance[:,14,...]
    Decoder2_logits[:,1,...] = pred[:,1,...]
    Decoder2_logits[:,2,...] = pred[:,2,...]
    Decoder2_logits[:,6,...] = pred[:,6,...]
    # Decoder_logits = Decoder1_logits+Decoder2_logits
    pred_chongtu = torch.argmax(Decoder2_logits, dim=1)
    # pred_pred = torch.argmax(Decoder1_logits, dim=1)
    pred_predinstance = torch.argmax(Decoder2_logits, dim=1)
    print(torch.unique(pred_chongtu))
    pred_chongtu[prediction_tmp == 255] = 0
    prediction_sum[prediction_tmp!=255] = 0
    prediction_final = (pred_chongtu + prediction_sum).cpu().numpy()
    prediction_to_save = label2color(prediction_final)[0]
    
    return prediction_to_save 


def Multi_box_prompts(input_prompt):
    with open('configs/fine_tuning_one_decoder.yaml', 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    model = models.make(config['model']).cpu()
    sam_checkpoint = torch.load("./save/model_epoch_last.pth", map_location='cpu')
    model.load_state_dict(sam_checkpoint, strict=False)
    model.eval()

    
    label2color = visual_utils.Label2Color(cmap=visual_utils.color_map('Unify_double'))
    # image_input.save(f'./save/visual_fair1m/input_img.png', quality=5)
    img = transforms.Resize([1024, 1024])(input_prompt["image"])
    input_img = transforms.ToTensor()(img).unsqueeze(0)
    image_embedding = model.image_encoder(input_img) # torch.Size([1, 256, 64, 64])
    sparse_embeddings, dense_embeddings, scatter_embeddings  = model.prompt_encoder(
        points=None,
        boxes=None,
        masks=None,
    scatter=None)
    # 目标类预测decoder
    low_res_masks, iou_predictions = model.mask_decoder(
        image_embeddings=image_embedding,
        image_pe=model.prompt_encoder.get_dense_pe(),
        sparse_prompt_embeddings=sparse_embeddings,
        dense_prompt_embeddings=dense_embeddings,
        multimask_output=False
    )
    pred = model.postprocess_masks(low_res_masks, model.inp_size, model.inp_size)
    _, prediction = pred.max(dim=1)
    prediction_to_save = label2color(prediction.cpu().numpy().astype(np.uint8))[0]

    def find_instance(image_map):
        BACKGROUND = 0
        steps = [[1, 0], [0, 1], [-1, 0], [0, -1], [1, 1], [1, -1], [-1, 1], [-1, -1]]
        instances = []

        def bfs(x, y, category_id):
            nonlocal image_map, steps
            instance = {(x, y)}
            q = [(x, y)]
            image_map[x, y] = BACKGROUND
            while len(q) > 0:
                x, y = q.pop(0)
                # print(x, y, image_map[x][y])
                for step in steps:
                    xx = step[0] + x
                    yy = step[1] + y
                    if 0 <= xx < len(image_map) and 0 <= yy < len(image_map[0]) \
                            and image_map[xx][yy] == category_id:  # and (xx, yy) not in q:
                        q.append((xx, yy))
                        instance.add((xx, yy))
                        image_map[xx, yy] = BACKGROUND
            return instance
        image_map = image_map[:]
        for i in range(len(image_map)):
            for j in range(len(image_map[i])):
                category_id = image_map[i][j]
                if category_id == BACKGROUND:
                    continue
                instances.append(bfs(i, j, category_id))
        return instances

    prompts = find_instance(np.uint8(np.array(input_prompt["mask"]).sum(-1) != 0))
    img_mask = np.array(img).copy()
    
    def get_box(prompt):
        xs = []
        ys = []
        for x, y in prompt:
            xs.append(x)
            ys.append(y)
        return [[min(xs), min(ys)], [max(xs), max(ys)]]

    def in_box(point, box):
        left_up, right_down = box
        x, y = point
        return x >= left_up[0] and x <= right_down[0] and y >= left_up[1] and y <= right_down[1]
        
    def draw_box(box_outer, img, radius=4):
        radius -= 1
        left_up_outer, right_down_outer = box_outer
        box_inner = [list(np.array(left_up_outer) + radius),
                                           list(np.array(right_down_outer) - radius)]
        for x in range(len(img)):
            for y in range(len(img[x])):
                if in_box([x, y], box_outer):
                    img_mask[x, y] = (1, 1, 1)
                if in_box([x, y], box_outer) and (not in_box([x, y], box_inner)):
                    img[x, y] = (255, 0, 0)
        return img
        
    for prompt in prompts:
        box = get_box(prompt)
        output = draw_box(box, prediction_to_save) * (img_mask==1)
    
    return output



def Decoder2_SAR(SAR_image, SAR_prompt):
    with open('configs/multi_mo_multi_task_sar_prompt.yaml', 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    model = models.make(config['model']).cpu()
    sam_checkpoint = torch.load("./save/SAR/model_epoch_last.pth", map_location='cpu')
    model.load_state_dict(sam_checkpoint, strict=True)
    model.eval()

    denorm = visual_utils.Denormalize(mean=[0.485, 0.456, 0.406],std=[0.229,0.224,0.225])
    label2color = visual_utils.Label2Color(cmap=visual_utils.color_map('Unify_YIJISAR'))

    img = transforms.Resize([1024, 1024])(SAR_image)
    transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229,0.224,0.225])])
    input_img = transform(img)
    input_img = torch.unsqueeze(input_img, dim=0)
    # input_img = transforms.ToTensor()(img).unsqueeze(0)
    # input_img = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229,0.224,0.225])
    filp_flag = torch.Tensor([False])
    image_embedding = model.image_encoder(input_img)
    
    # scattter_prompt = cv2.imread(scatter_file_, cv2.IMREAD_UNCHANGED)
    # scattter_prompt = get_prompt_inp_scatter(name[0].replace('gt', 'JIHUAFENJIE'))
    SAR_prompt = cv2.imread(SAR_prompt, cv2.IMREAD_UNCHANGED)
    scatter_torch = pre_scatter_prompt(SAR_prompt, filp_flag, device=input_img.device)
    scatter_torch = scatter_torch.unsqueeze(0)
    scatter_torch = torch.nn.functional.interpolate(scatter_torch, size=(256, 256))
    sparse_embeddings, dense_embeddings, scatter_embeddings = model.prompt_encoder(
        points=None,
        boxes=None,
        masks=None,
        scatter=scatter_torch)
    # 地物类预测decoder
    low_res_masks, iou_predictions_2 = model.mask_decoder_diwu(
        image_embeddings=image_embedding,
        image_pe=model.prompt_encoder.get_dense_pe(),
        sparse_prompt_embeddings=sparse_embeddings,
        dense_prompt_embeddings=dense_embeddings,
        # multimask_output=False,
        multimask_output=True,
    )  # B*C+1*H*W
    pred = model.postprocess_masks(low_res_masks, model.inp_size, model.inp_size)
    _, prediction = pred.max(dim=1)
    prediction = prediction.cpu().numpy()
    prediction_to_save = label2color(prediction)[0]
    
    return prediction_to_save


examples1_instance = [
            ['./images/optical/isaid/_P0007_1065_319_image.png'], 
            ['./images/optical/isaid/_P0466_1068_420_image.png'], 
            ['./images/optical/isaid/_P0897_146_34_image.png'],
            ['./images/optical/isaid/_P1397_844_904_image.png'], 
            ['./images/optical/isaid/_P2645_883_965_image.png'],
            ['./images/optical/isaid/_P1398_1290_630_image.png']
           ]

examples1_terrain = [
            ['./images/optical/vaihingen/top_mosaic_09cm_area2_105_image.png'],
            ['./images/optical/vaihingen/top_mosaic_09cm_area4_227_image.png'],
            ['./images/optical/vaihingen/top_mosaic_09cm_area20_142_image.png'],
            ['./images/optical/vaihingen/top_mosaic_09cm_area24_128_image.png'],
            ['./images/optical/vaihingen/top_mosaic_09cm_area27_34_image.png']
           ]


examples1_multi_box = [
            ['./images/optical/isaid/_P0007_1065_319_image.png'], 
            ['./images/optical/isaid/_P0466_1068_420_image.png'], 
            ['./images/optical/isaid/_P0897_146_34_image.png'],
            ['./images/optical/isaid/_P1397_844_904_image.png'], 
            ['./images/optical/isaid/_P2645_883_965_image.png'],
            ['./images/optical/isaid/_P1398_1290_630_image.png']
           ]


examples2 = [
            ['./images/sar/YIJISARGF3_MYN_QPSI_001269_E113.2_N23.0_20161105_L1A_L10002009158_ampl_4_image.png', './images/sar/YIJISARGF3_MYN_QPSI_001269_E113.2_N23.0_20161105_L1A_L10002009158_ampl_4.png'],
            ['./images/sar/YIJISARGF3_MYN_QPSI_001269_E113.2_N23.0_20161105_L1A_L10002009158_ampl_15_image.png', './images/sar/YIJISARGF3_MYN_QPSI_001269_E113.2_N23.0_20161105_L1A_L10002009158_ampl_15.png'],
            ['./images/sar/YIJISARGF3_MYN_QPSI_001269_E113.2_N23.0_20161105_L1A_L10002009158_ampl_24_image.png', './images/sar/YIJISARGF3_MYN_QPSI_001269_E113.2_N23.0_20161105_L1A_L10002009158_ampl_24.png'],
            ['./images/sar/YIJISARGF3_MYN_QPSI_001269_E113.2_N23.0_20161105_L1A_L10002009158_ampl_41_image.png', './images/sar/YIJISARGF3_MYN_QPSI_001269_E113.2_N23.0_20161105_L1A_L10002009158_ampl_41.png'],
            ['./images/sar/YIJISARGF3_MYN_QPSI_999996_E121.2_N30.3_20160815_L1A_L10002015572_ampl_150_image.png', './images/sar/YIJISARGF3_MYN_QPSI_999996_E121.2_N30.3_20160815_L1A_L10002015572_ampl_150.png']
           ]



# RingMo-SAM designs two new promptable forms based on the characteristics of multimodal remote sensing images: 
# multi-boxes prompt and SAR polarization scatter prompt.


title = "RingMo-SAM:A Foundation Model for Segment Anything in Multimodal Remote Sensing Images<br> \
<div align='center'> \
<h2><a href='https://ieeexplore.ieee.org/document/10315957' target='_blank' rel='noopener'>[paper]</a> \
<br> \
<image src='file/RingMo-SAM.gif' width='720px' /> \
<h2>RingMo-SAM can not only segment anything in optical and SAR remote sensing data, but also identify object categories.<h2> \
</div> \
"

# <a href='https://github.com/AICyberTeam' target='_blank' rel='noopener'>[code]</a></h2> \
# with gr.Blocks() as demo:
#     image_input = gr.Image(type='pil', label='Input Img')
#     image_output = gr.Image(label='Segment Result', type='numpy')
    
        
Decoder_optical_instance_io = gr.Interface(fn=Decoder1_optical_instance,
                  inputs=[gr.Image(type='pil', label='optical_instance_img(光学图像)')],
                  outputs=[gr.Image(label='segment_result', type='numpy')],
                  # title=title,
                  description="<p> \
                    Instance_Decoder:<br>\
                    Instance-type objects (such as vehicle, aircraft, ship, etc.) have a smaller proportion. <br>\
                    Our decoder can decouple the SAM's mask decoder into instance category decoder and terrain category decoder to ensure that the model fits adequately to both types of data. <br>\
                    Choose an example below, or, upload optical instance images to be tested. <br>\
                    Examples below were never trained and are randomly selected for testing in the wild. <br>\
                  </p>",
                  allow_flagging='auto',
                  examples=examples1_instance,
                  cache_examples=False,
                  )


Decoder_optical_terrain_io = gr.Interface(fn=Decoder1_optical_terrain,
                  inputs=[gr.Image(type='pil', label='optical_terrain_img(光学图像)')],
                  # inputs=[gr.Image(type='pil', label='optical_img(光学图像)'), gr.Image(type='pil', label='SAR_img(SAR图像)'), gr.Image(type='pil', label='SAR_prompt(偏振散射提示)')],
                  outputs=[gr.Image(label='segment_result', type='numpy')],
                  # title=title,
                  description="<p> \
                    Terrain_Decoder:<br>\
                    Terrain-type objects (such as vegetation, land, river, etc.) have a larger proportion. <br>\
                    Our decoder can decouple the SAM's mask decoder into instance category decoder and terrain category decoder to ensure that the model fits adequately to both types of data. <br>\
                    Choose an example below, or, upload optical terrain images to be tested. <br>\
                    Examples below were never trained and are randomly selected for testing in the wild. <br>\
                  </p>",
                  allow_flagging='auto',
                  examples=examples1_terrain,
                  cache_examples=False,
                  )



Decoder_multi_box_prompts_io = gr.Interface(fn=Multi_box_prompts,
              inputs=[gr.ImageMask(brush_radius=4, type='pil', label='input_img(图像)')],
              outputs=[gr.Image(label='segment_result', type='numpy')],
              # title=title,
              description="<p> \
                Multi-box Prompts:<br>\
                Multiple boxes are sequentially encoded as concated sparse high-dimensional feature embedding, \
                the corresponding multiple high-dimensional features are concated together into a high-dimensional feature vector as part of the sparse embedding. <br>\
                Choose an example below, or, upload images to be tested, and then draw multi-boxes.  <br>\
                Examples below were never trained and are randomly selected for testing in the wild. <br>\
              </p>",
              allow_flagging='auto',
              examples=examples1_multi_box,
              cache_examples=False,
              )



Decoder_SAR_io = gr.Interface(fn=Decoder2_SAR,
                  inputs=[gr.Image(type='pil', label='SAR_img(SAR图像)'), gr.Image(type='filepath', label='SAR_prompt(偏振散射提示)')],
                  outputs=[gr.Image(label='segment_result', type='numpy')],
                  description="<p> \
                    SAR Polarization Scatter Prompts:<br>\
                    Different terrain categories usually exhibit different scattering properties. \
                    Therefore, we code network for coded mapping of these SAR polarization scatter prompts to the corresponding SAR images, \
                    which improves the segmentation results of SAR images. <br>\
                    Choose an example below, or, upload SAR images and the corresponding polarization scatter prompts to be tested. <br>\
                    Examples below were never trained and are randomly selected for testing in the wild. <br>\
                  </p>",           
                  allow_flagging='auto',
                  examples=examples2,
                  cache_examples=False,
                  )


# Decoder1_io.launch(server_name="0.0.0.0", server_port=34311)
# Decoder1_io.launch(enable_queue=False)
# demo = gr.TabbedInterface([Decoder1_io, Decoder2_io], ['Instance_Decoder', 'Terrain_Decoder'], title=title)
demo = gr.TabbedInterface([Decoder_optical_instance_io, Decoder_optical_terrain_io, Decoder_multi_box_prompts_io, Decoder_SAR_io], ['optical_instance_img(光学图像)', 'optical_terrain_img(光学图像)', 'multi_box_prompts(多框提示)', 'SAR_img(偏振散射提示)'], title=title).launch()
# -