import torch import gradio as gr from PIL import Image import cv2 from AV.models.network import PGNet from AV.Tools.AVclassifiation import AVclassifiation from AV.Tools.utils_test import paint_border_overlap, extract_ordered_overlap_big, Normalize, sigmoid, recompone_overlap, \ kill_border from AV.config import config_test_general as cfg import torch.autograd as autograd import numpy as np import os from datetime import datetime from huggingface_hub import hf_hub_download hf_token = os.environ.get("HF_token") def creatMask(Image, threshold=5): ##This program try to creat the mask for the filed-of-view ##Input original image (RGB or green channel), threshold (user set parameter, default 10) ##Output: the filed-of-view mask if len(Image.shape) == 3: ##RGB image gray = cv2.cvtColor(Image, cv2.COLOR_BGR2GRAY) Mask0 = gray >= threshold else: # for green channel image Mask0 = Image >= threshold # ######get the largest blob, this takes 0.18s cvVersion = int(cv2.__version__.split('.')[0]) Mask0 = np.uint8(Mask0) contours, hierarchy = cv2.findContours(Mask0, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) areas = [cv2.contourArea(c) for c in contours] max_index = np.argmax(areas) Mask = np.zeros(Image.shape[:2], dtype=np.uint8) cv2.drawContours(Mask, contours, max_index, 1, -1) ResultImg = Image.copy() if len(Image.shape) == 3: ResultImg[Mask == 0] = (255, 255, 255) else: ResultImg[Mask == 0] = 255 Mask[Mask > 0] = 255 kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3)) Mask = cv2.morphologyEx(Mask, cv2.MORPH_OPEN, kernel, iterations=3) return ResultImg, Mask def shift_rgb(img, *args): result_img = np.empty_like(img) shifts = args max_value = 255 # print(shifts) for i, shift in enumerate(shifts): lut = np.arange(0, max_value + 1).astype("float32") lut += shift lut = np.clip(lut, 0, max_value).astype(img.dtype) if len(img.shape) == 2: print(f'=========grey image=======') result_img = cv2.LUT(img, lut) else: result_img[..., i] = cv2.LUT(img[..., i], lut) return result_img def CAM(x, img_path, rate=0.8, ind=0): """ :param dataset_path: 计算整个训练数据集的平均RGB通道值 :param image: array， 单张图片的array 形式 :return: array形式的cam后的结果 """ # 每次使用新数据集时都需要重新计算前面的RBG平均值 # RGB-->Rshift-->CLAHE x = np.uint8(x) _, Mask0 = creatMask(x, threshold=10) Mask = np.zeros((x.shape[0], x.shape[1]), np.float32) Mask[Mask0 > 0] = 1 resize = False R_mea_num, G_mea_num, B_mea_num = [], [], [] dataset_path = img_path image = np.array(Image.open(dataset_path)) R_mea_num.append(np.mean(image[:, :, 0])) G_mea_num.append(np.mean(image[:, :, 1])) B_mea_num.append(np.mean(image[:, :, 2])) mea2stand = int((np.mean(R_mea_num) - np.mean(x[:, :, 0])) * rate) mea2standg = int((np.mean(G_mea_num) - np.mean(x[:, :, 1])) * rate) mea2standb = int((np.mean(B_mea_num) - np.mean(x[:, :, 2])) * rate) y = shift_rgb(x, mea2stand, mea2standg, mea2standb) y[Mask == 0, :] = 0 return y def modelEvalution_out_big(net, use_cuda=False, dataset='', is_kill_border=True, input_ch=3, config=None, output_dir='', evaluate_metrics=False): # path for images to save n_classes = 3 Net = PGNet(use_global_semantic=config.use_global_semantic, input_ch=input_ch, num_classes=n_classes, use_cuda=use_cuda, pretrained=False, centerness=config.use_centerness, centerness_map_size=config.centerness_map_size) msg = Net.load_state_dict(net, strict=False) if use_cuda: Net.cuda() Net.eval() image_basename = dataset # if not os.path.exists(output_dir): # os.makedirs(output_dir) step = 1 # every step of between star and end for loop until len(image_basename) # for start_end in start_end_list: image0 = cv2.imread(image_basename) test_image_height = image0.shape[0] test_image_width = image0.shape[1] if config.use_resize: if min(test_image_height, test_image_width) <= 256: scaling = 512 / min(test_image_height, test_image_width) new_width = int(test_image_width * scaling) new_height = int(test_image_height * scaling) test_image_width, test_image_height = new_width, new_height # 大尺寸处理：确保最长边≤1536 elif max(test_image_height, test_image_width) >= 2048: scaling = 2048 / max(test_image_height, test_image_width) new_width = int(test_image_width * scaling) new_height = int(test_image_height * scaling) test_image_width, test_image_height = new_width, new_height ArteryPredAll = np.zeros((1, 1, test_image_height, test_image_width), np.float32) VeinPredAll = np.zeros((1, 1, test_image_height, test_image_width), np.float32) VesselPredAll = np.zeros((1, 1, test_image_height, test_image_width), np.float32) ProMap = np.zeros((1, 3, test_image_height, test_image_width), np.float32) MaskAll = np.zeros((1, 1, test_image_height, test_image_width), np.float32) ArteryPred, VeinPred, VesselPred, Mask, LabelArtery, LabelVein, LabelVessel = GetResult_out_big(Net, 0, use_cuda=use_cuda, dataset=image_basename, is_kill_border=is_kill_border, config=config, resize_w_h=( test_image_width, test_image_height) ) ArteryPredAll[0 % step, :, :, :] = ArteryPred VeinPredAll[0 % step, :, :, :] = VeinPred VesselPredAll[0 % step, :, :, :] = VesselPred MaskAll[0 % step, :, :, :] = Mask image_color = AVclassifiation(output_dir, ArteryPredAll, VeinPredAll, VesselPredAll, 1, image_basename) return image_color def GetResult_out_big(Net, k, use_cuda=False, dataset='', is_kill_border=False, config=None, resize_w_h=None): ImgName = dataset Img0 = cv2.imread(ImgName) _, Mask0 = creatMask(Img0, threshold=-1) Mask = np.zeros((Img0.shape[0], Img0.shape[1]), np.float32) Mask[Mask0 > 0] = 1 if config.use_resize: Img0 = cv2.resize(Img0, resize_w_h) Mask = cv2.resize(Mask, resize_w_h, interpolation=cv2.INTER_NEAREST) Img = Img0 height, width = Img.shape[:2] n_classes = 3 patch_height = config.patch_size patch_width = config.patch_size stride_height = config.stride_height stride_width = config.stride_width Img = cv2.cvtColor(Img, cv2.COLOR_BGR2RGB) if cfg.dataset == 'all': # # # 将图像转换为 LAB 颜色空间 lab = cv2.cvtColor(Img, cv2.COLOR_RGB2LAB) # 拆分 LAB 通道 l, a, b = cv2.split(lab) # 创建 CLAHE 对象并应用到 L 通道 clahe = cv2.createCLAHE(clipLimit=2, tileGridSize=(8, 8)) l_clahe = clahe.apply(l) # 将 CLAHE 处理后的 L 通道与原始的 A 和 B 通道合并 lab_clahe = cv2.merge((l_clahe, a, b)) # 将图像转换回 BGR 颜色空间 Img = cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2RGB) if cfg.use_CAM: Img = CAM(Img, dataset) Img = np.float32(Img / 255.) Img_enlarged = paint_border_overlap(Img, patch_height, patch_width, stride_height, stride_width) patch_size = config.patch_size batch_size = 2 patches_imgs, global_images = extract_ordered_overlap_big(Img_enlarged, patch_height, patch_width, stride_height, stride_width) patches_imgs = np.transpose(patches_imgs, (0, 3, 1, 2)) patches_imgs = Normalize(patches_imgs) global_images = np.transpose(global_images, (0, 3, 1, 2)) global_images = Normalize(global_images) patchNum = patches_imgs.shape[0] max_iter = int(np.ceil(patchNum / float(batch_size))) pred_patches = np.zeros((patchNum, n_classes, patch_size, patch_size), np.float32) for i in range(max_iter): begin_index = i * batch_size end_index = (i + 1) * batch_size patches_temp1 = patches_imgs[begin_index:end_index, :, :, :] patches_input_temp1 = torch.FloatTensor(patches_temp1) global_input_temp1 = patches_input_temp1 if config.use_global_semantic: global_temp1 = global_images[begin_index:end_index, :, :, :] global_input_temp1 = torch.FloatTensor(global_temp1) if use_cuda: patches_input_temp1 = autograd.Variable(patches_input_temp1.cuda()) if config.use_global_semantic: global_input_temp1 = autograd.Variable(global_input_temp1.cuda()) else: patches_input_temp1 = autograd.Variable(patches_input_temp1) if config.use_global_semantic: global_input_temp1 = autograd.Variable(global_input_temp1) output_temp, _1, = Net(patches_input_temp1, global_input_temp1) pred_patches_temp = np.float32(output_temp.data.cpu().numpy()) pred_patches_temp_sigmoid = sigmoid(pred_patches_temp) pred_patches[begin_index:end_index, :, :, :] = pred_patches_temp_sigmoid[:, :, :patch_size, :patch_size] del patches_input_temp1 del pred_patches_temp del patches_temp1 del output_temp del pred_patches_temp_sigmoid new_height, new_width = Img_enlarged.shape[0], Img_enlarged.shape[1] pred_img = recompone_overlap(pred_patches, new_height, new_width, stride_height, stride_width) # predictions pred_img = pred_img[:, 0:height, 0:width] if is_kill_border: pred_img = kill_border(pred_img, Mask) ArteryPred = np.float32(pred_img[0, :, :]) VeinPred = np.float32(pred_img[2, :, :]) VesselPred = np.float32(pred_img[1, :, :]) ArteryPred = ArteryPred[np.newaxis, :, :] VeinPred = VeinPred[np.newaxis, :, :] VesselPred = VesselPred[np.newaxis, :, :] Mask = Mask[np.newaxis, :, :] return ArteryPred, VeinPred, VesselPred, Mask, ArteryPred, VeinPred, VesselPred, def out_test(cfg,model_path='', output_dir='', evaluate_metrics=False, img_name='out_test'): device = torch.device("cuda" if cfg.use_cuda else "cpu") model_path = model_path net = torch.load(model_path, map_location=device) image_color = modelEvalution_out_big(net, use_cuda=cfg.use_cuda, dataset=img_name, input_ch=cfg.input_nc, config=cfg, output_dir=output_dir, evaluate_metrics=evaluate_metrics) return image_color def segment_by_out_test(image,model_name): print("✅ 传到后端的模型名：", model_name) model_path = hf_hub_download( repo_id="weidai00/RIP-AV-sulab", # 模型库的名字 filename=f"G_{model_name}.pkl", # 文件名 repo_type="model", # 模型库必须写 repo_type token=hf_token ) cfg.set_dataset(model_name) if image is None: raise gr.Error("请上传一张图像(upload a fundus image)。") os.makedirs("./examples", exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") temp_path = f"./examples/tmp_upload_{timestamp}.png" image.save(temp_path) image_color = out_test(cfg,model_path=model_path, output_dir='', evaluate_metrics=False, img_name=temp_path) return Image.fromarray(image_color) def gradio_interface(): model_info_md = """ ### 📘 模型说明 | 模型(model name) | 数据集(dataset) | patch size |running time | |------|--------|------------|--------| | DRIVE | 小分辨率血管图像 | 256 |30s以内| | HRF | 高分辨率图像（健康、青光眼等）| 256 | 2min以内| | LES | 视盘中心图像适配 | 256 |2min以内| | UKBB | UKBB图像 | 256 |2min以内 | | 通用模型（512） | 超清图像，适配性强 | 512 |2min以内| """ model_choices = [ ("1: DRIVE专用模型", "DRIVE"), ("2: HRF专用模型", "hrf"), ("3: LES专用模型","LES"), ("4: UKBB专用模型", "ukbb"), ("5: 通用模型(general)", "all"), ] with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# 👁️ 眼底图像动静脉血管分割(Retinal image artery and vein segmentation)") gr.Markdown("上传眼底图像，选择一个模型开始处理，结果将自动生成。(Upload the retinal image, select a model to start processing, and the results will be generated automatically.)") with gr.Row(): image_input = gr.Image(type="pil", label="📤 上传图像(upload)",height=300) with gr.Row(): with gr.Column(): model_select = gr.Radio( choices=model_choices, label="🎯 选择模型", value="DRIVE", interactive = True ) submit_btn = gr.Button("🚀 开始分割(RUN)") with gr.Column(): output_image = gr.Image(label="🖼️ 分割结果(Result)") gr.Markdown("### 📁 示例图像examples（点击自动加载）") gr.Examples( examples=[ ["examples/DRIVE.tif", "DRIVE"], ["examples/LES.png", "LES"], ["examples/hrf.png", "hrf"], ["examples/ukbb.png", "ukbb"], ["examples/all.jpg", "all"] ], inputs=[image_input, model_select], label="示例图像", examples_per_page=5 ) with gr.Accordion("📖 模型说明-Description（点击展开）", open=False): gr.Markdown(model_info_md) # 功能连接 submit_btn.click( fn=segment_by_out_test, inputs=[image_input, model_select], outputs=[output_image] ) gr.Markdown("📚 **专用模型引用cite**: RIP-AV: Joint Representative Instance Pre-training with Context Aware Network for Retinal Artery/Vein Segmentation") gr.Markdown("📚 **通用模型引用cite**: An Efficient and Interpretable Foundation Model for Retinal Image Analysis in Disease Diagnosis.") demo.queue() demo.launch() if __name__ == '__main__': # cfg.set_dataset('all') # image_color = out_test(cfg = cfg, evaluate_metrics=False, img_name=r'.\AV\data\AV-DRIVE\test\images\01_test.tif') # Image.fromarray(image_color).save('image_color.png') #print(cfg.patch_size) gradio_interface()