Upload 4 files

UTILS/__init__.py

@@ -0,0 +1 @@
+from .more_dim import get_more_dim

UTILS/more_dim.py

@@ -0,0 +1,93 @@
+import traceback
+
+import numpy as np
+import cv2 as cv
+from matplotlib import pyplot as plt
+
+
+def get_binary_img_(img):
+    gray_img = img
+    if len(img.shape) > 2:
+        gray_img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
+    binary_img = cv.Canny(gray_img, 80, 150)
+    return binary_img
+
+
+def get_morp_dilate_(binary_img):
+    kernel = cv.getStructuringElement(cv.MORPH_RECT, (3, 3))
+    # morp_dilate = cv.morphologyEx(binaryImg, cv.MORPH_DILATE, kernel=(1, 3), iterations=3)
+    # morp_dilate = cv.morphologyEx(morp_dilate, cv.MORPH_DILATE, kernel=(3, 1), iterations=3)
+    # morp_dilate = cv.morphologyEx(binaryImg, cv.MORPH_DILATE, kernel=(11, 11), iterations=3)
+    morp_dilate = cv.morphologyEx(binary_img, cv.MORPH_DILATE, kernel=kernel, iterations=3)
+    return morp_dilate
+
+
+def get_water_img_(img, morp_dilate):
+    # 寻找图像轮廓 返回修改后的 图像的轮廓  以及它们的层次
+    # contours, hierarchy = cv.findContours(gray_img, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
+    # contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
+    # contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
+    contours, hierarchy = cv.findContours(morp_dilate, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
+    # 32位有符号整数类型，
+    marks = np.zeros(morp_dilate.shape[:2], np.int32)
+    # 绘制每一个轮廓
+    for index in range(len(contours)):
+        # 对marks进行标记，对不同区域的轮廓使用不同的亮度绘制，相当于设置注水点，有多少个轮廓，就有多少个轮廓
+        # 图像上不同线条的灰度值是不同的，底部略暗，越往上灰度越高
+        marks = cv.drawContours(marks, contours, index, (index, index, index), 1, 8, hierarchy)
+
+    # 使用分水岭算法
+    # 经过watershed函数的处理，不同区域间的值被置为-1（边界）没有标记清楚的区域被置为0，其他每个区域的值保持不变：1，2，...，contours.size()
+    marks_water = cv.watershed(img, marks)
+    return marks_water
+
+
+def get_mask_img_(morp_dilate, file_dir):
+    contours, hierarchy = cv.findContours(morp_dilate, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
+    # 32位有符号整数类型，
+    marks = np.zeros(morp_dilate.shape[:2], np.int32)
+    for index in range(len(contours)):
+        dist = cv.pointPolygonTest(contours[index], (marks.shape[0] // 2, marks.shape[1] // 2), True)
+        if dist >= 0:
+            marks = cv.drawContours(marks, contours, contourIdx=index, color=1, thickness=1, lineType=8,
+                                    hierarchy=hierarchy)
+
+    edges = np.zeros((marks.shape[0] + 2, marks.shape[1] + 2), np.uint8)  # 掩码，长短需要加2个像素
+    try:
+        cv.floodFill(marks, edges, (marks.shape[0] // 2, marks.shape[1] // 2), 1, cv.FLOODFILL_MASK_ONLY)  # 漫水填充
+    except Exception as e:
+        if file_dir:
+            print(file_dir)
+            print(e)
+            print("=================")
+            print(traceback.format_exc())
+        # raise e
+        marks = np.ones(morp_dilate.shape[:2], np.int32)
+    return marks
+
+
+def get_binary_img(binary_img, mask):
+    masked_binary_img = cv.bitwise_and(binary_img, binary_img, mask=mask.astype('uint8'))
+    return masked_binary_img
+
+
+def get_water_img(img, morp_dilate, mask):
+    water_img = get_water_img_(img, morp_dilate)
+    masked_water = cv.bitwise_and(water_img, water_img, mask=mask.astype('uint8'))
+    return masked_water
+
+
+def get_more_dim(img, file_dir, source_img=None):
+    if source_img is None:
+        source_img = img
+    # img: ndarray: 852, 847, 3
+    binary_img = get_binary_img_(img)
+    morp_dilate = get_morp_dilate_(binary_img)
+    mask = get_mask_img_(morp_dilate, file_dir)
+
+    masked_binary_img = get_binary_img(binary_img, mask)
+    masked_water = get_water_img(source_img, morp_dilate, mask)
+    # print(f"masked_binary_img shape:{masked_binary_img.shape} masked_water shape:{masked_water.shape}")
+    # print(f"type(masked_binary_img):{type(masked_binary_img)} type(masked_water):{type(masked_water)}")
+    # return np.stack((masked_binary_img, mask), axis=0)
+    return masked_binary_img, masked_water

app.py

@@ -0,0 +1,181 @@
+import logging
+import traceback
+from io import BytesIO
+
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+from matplotlib import pyplot as plt
+from matplotlib.colors import ListedColormap
+from torchvision import transforms
+
+from UTILS import get_more_dim
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model_path = 'models/model_v48.pth'
+model_pic_size = 512
+model_class_num = 14
+model = torch.load(model_path, map_location=torch.device('cpu'))
+model = model.to(device)
+
+colors = ['Black', 'Silver', 'White', 'Brown', 'LightCoral', 'Tomato', 'LightSalmon', 'Chocolate', 'Tan',
+          'PapayaWhip', 'Gold', 'Ivory', 'GreenYellow', 'Green', 'DarkSeaGreen', 'DarkTurquoise', 'LightBLue',
+          'SteelBlue']
+mode = 'predict'
+
+
+def get_predict(origin_img, need_subplot=False):
+    features, pad_width, pad_height = get_features(origin_img, pic_size=model_pic_size)
+    predict_npy, subplot_img = save_predict(model, features, device=device, class_num=model_class_num,
+                                            need_subplot=need_subplot)
+    return predict_npy, subplot_img, pad_width, pad_height
+
+
+def save_predict(model, features, device, class_num=14, need_subplot=False):
+    cmap = ListedColormap(colors[:class_num])
+
+    model.eval()
+    with torch.no_grad():
+        features = features.to(device)
+
+        predictions = model(features)
+
+        features = torch.squeeze(features)
+        features = features.detach().cpu()
+        predictions = torch.squeeze(predictions)
+        predictions = predictions.detach().cpu()
+
+        features_len = features.shape[0]
+
+        origin_img = transforms.ToPILImage()(features[:3])
+        binary_img = features[3]
+        water_img = features[4]
+        predict_img = label_to_img(predictions)
+        predict_npy = predict_img.numpy().astype('uint8')
+
+        subplot = None
+        if need_subplot:
+            subplot = save_subplot(features_len, origin_img, predict_img, binary_img, water_img, vmax=class_num,
+                                   cmap=cmap)
+
+        return predict_npy, subplot
+
+
+def label_to_img(label):
+    max_label_values, max_label_indices = torch.max(label, dim=0)
+    return max_label_indices
+
+
+def save_subplot(features_len, origin_img, predict_img, feature_1=None, feature_2=None, vmax=14,
+                 cmap=None):
+    plt.clf()
+    plt.close()
+
+    # colorbar 左 下 宽 高 ；设置colorbar位置；
+    rect = [0.92, 0.36, 0.015, 0.99 - 0.37 * 2]
+
+    fig = plt.figure()
+    subplot_num = features_len - 2 + 1
+    subplot_count = 0
+
+    subplot_count += 1
+    plt.subplot(1, subplot_num, subplot_count)
+    plt.imshow(origin_img)
+    if features_len > 3:
+        subplot_count += 1
+        plt.subplot(1, subplot_num, subplot_count)
+        plt.imshow(feature_1)
+    if features_len > 4:
+        subplot_count += 1
+        plt.subplot(1, subplot_num, subplot_count)
+        plt.imshow(feature_2)
+
+    subplot_count += 1
+    plt.subplot(1, subplot_num, subplot_count)
+    im = plt.imshow(predict_img, vmin=-1, vmax=vmax, cmap=cmap)
+    # 前面三个子图的总宽度 为 全部宽度的 0.9；剩下的0.1用来放置colorbar
+    fig.subplots_adjust(right=0.9)
+    cbar_ax = fig.add_axes(rect)
+    plt.colorbar(im, cax=cbar_ax)
+
+    with BytesIO() as out:
+        plt.savefig(out, dpi=300)
+        subplot_bytes = out.getvalue()
+    return subplot_bytes
+
+
+def get_features(origin_img, pic_size):
+    img = origin_img.convert('RGB')
+    img_np = np.array(img)
+    try:
+        masked_binary_img, masked_water = get_more_dim(img_np, file_dir=None)
+    except Exception as e:
+        logging.error(e)
+        logging.error("=================")
+        logging.error(traceback.format_exc())
+        masked_binary_img = np.zeros(img_np.shape[:2], np.int32)
+        masked_water = np.zeros(img_np.shape[:2], np.int32)
+    img, pad_width, pad_height = transform_pic_shape(img, pic_size)
+    masked_binary_img, _, _ = transform_pic_shape(torch.tensor(masked_binary_img), pic_size)
+    masked_water, _, _ = transform_pic_shape(torch.tensor(masked_water), pic_size)
+    data_mode_dim = torch.stack((masked_binary_img, masked_water), axis=0)
+    img = transforms.ToTensor()(img)
+    featurs = torch.cat((img, data_mode_dim), dim=0)
+    featurs = torch.unsqueeze(featurs, dim=0)
+    return featurs, pad_width, pad_height
+
+
+def transform_pic_shape(img, pic_size):
+    # 对于RGB图
+    # Image.size为（宽，高）
+    # array.shape为（高，宽，通道数）
+    # array.size为 高x宽x通道数 的总个数
+    height, width = get_image_shape(img)
+    if height > pic_size - 1 or width > pic_size - 1:
+        is_unsqueeze = False
+        if type(img) == torch.Tensor and len(img.shape) == 2:
+            img = torch.unsqueeze(img, dim=0)
+            is_unsqueeze = True
+        img = transforms.Resize(size=pic_size - 1, max_size=pic_size,
+                                interpolation=transforms.InterpolationMode.NEAREST)(img)
+        if is_unsqueeze:
+            img = torch.squeeze(img)
+        height, width = get_image_shape(img)
+
+    pad_width = 0
+    pad_height = 0
+    if height < pic_size or width < pic_size:
+        # 当为 a 时，上下左右均填充 a 个像素
+        # 当为 (a, b) 时，左右填充 a 个像素，上下填充 b 个像素
+        # 当为 (a, b, c, d) 时，左上右下分别填充 a，b，c，d
+        # padding_mode: 填充模式，有 4 种模式，constant、edge、reflect、symmetric
+        pad_width = (pic_size - width) // 2
+        pad_height = (pic_size - height) // 2
+        img = transforms.Pad(
+            padding=[pad_width, pad_height, pic_size - pad_width - width, pic_size - pad_height - height],
+            fill=0)(img)
+    return img, pad_width, pad_height
+
+
+def get_image_shape(img):
+    if type(img) == Image.Image:
+        width, height = img.size
+    else:
+        if len(img.shape) == 3:
+            channel_num, height, width = img.shape
+        else:
+            height, width = img.shape
+    return height, width
+
+
+def greet(img):
+    predict_npy, subplot_img, pad_width, pad_height = get_predict(img, need_subplot=False)
+    predict_npy = predict_npy / model_class_num * 255
+    predict_img = Image.fromarray(predict_npy).convert(mode='L')
+    return predict_img
+
+
+iface = gr.Interface(fn=greet, inputs=gr.Image(type="pil"), outputs="image")
+iface.launch(server_name="0.0.0.0", share=True)

models/model_v48.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1cd93eeccf176fcee883213eef18edfc2098d40733279617b43c070ae73227c9
+size 183422311