atom-predict/msunet/pre.py

119 lines
3.3 KiB
Python
Executable File

import os
import cv2
import glob
import copy
import json
import shutil
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from PIL import Image
from labelme import utils
from skimage.feature import peak_local_max
class_dict = {
1: 'Norm',
2: 'SV',
3: 'LineSV',
}
# class_dict_rev = {
# '0': 1,
# '1': 2,
# '2': 3,
# }
class_dict_rev = {
'Norm': 1,
'1': 2,
'2': 3,
}
def crop_slide(img_path, save_path, patch_size=256, step=128):
# 2048*2048,裁减出来的size是256*256,滑动窗口128
base_name = img_path.split('/')[-1].split('.')[0]
json_path = img_path.replace('.jpg', '.json')
img = cv2.imread(img_path, 0)
# img = cv2.equalizeHist(img)
# img = cv2.GaussianBlur(img, (5, 5), 0)
h, w = img.shape
with open(json_path) as f:
json_data = json.load(f)
points = np.array([item['points'][0][::-1] for item in json_data['shapes']], np.int32)
labels = np.array([class_dict_rev[item['label']] for item in json_data['shapes']], np.int32)
mask = np.zeros_like(img)
mask[points[:, 0], points[:, 1]] = labels
img_item_path = os.path.join(save_path, 'img')
if not os.path.exists(img_item_path):
os.makedirs(img_item_path, exist_ok=True)
lbl_item_path = os.path.join(save_path, 'lbl')
if not os.path.exists(lbl_item_path):
os.makedirs(lbl_item_path, exist_ok=True)
for i in range(0, h - patch_size + 1, step):
for j in range(0, w - patch_size + 1, step):
v_nums = np.sum(mask[i:i + patch_size, j:j + patch_size] > 1)
Image.fromarray(img[i:i + patch_size, j:j + patch_size]).save(
os.path.join(save_path, 'img', '{}_{}_{}_{}.png'.format(base_name, str(i), str(j), str(v_nums)))
)
Image.fromarray(mask[i:i + patch_size, j:j + patch_size]).save(
os.path.join(save_path, 'lbl', '{}_{}_{}_{}.png'.format(base_name, str(i), str(j), str(v_nums)))
)
def process_slide(img_path, save_path):
base_name = img_path.split('/')[-1].split('.')[0]
json_path = img_path.replace('.jpg', '.json')
img = cv2.imread(img_path, 0)
# img = cv2.equalizeHist(img)
# img = cv2.GaussianBlur(img, (5, 5), 0)
# h, w = img.shape
with open(json_path) as f:
json_data = json.load(f)
points = np.array([item['points'][0][::-1] for item in json_data['shapes']], np.int32)
labels = np.array([class_dict_rev[item['label']] for item in json_data['shapes']], np.int32)
# mask = np.zeros_like(img)
# for idx, point in enumerate(points):
# cv2.circle(mask, point[::-1], 8, int(labels[idx]), -1)
mask = np.zeros_like(img)
mask[points[:, 0], points[:, 1]] = labels
img_item_path = os.path.join(save_path, 'img')
if not os.path.exists(img_item_path):
os.makedirs(img_item_path, exist_ok=True)
lbl_item_path = os.path.join(save_path, 'lbl')
if not os.path.exists(lbl_item_path):
os.makedirs(lbl_item_path, exist_ok=True)
Image.fromarray(img).save(
os.path.join(save_path, 'img', '{}.png'.format(base_name))
)
Image.fromarray(mask).save(
os.path.join(save_path, 'lbl', '{}.png'.format(base_name))
)
# img_lst = glob.glob('./train_and_test/test/*.jpg')
# for item in img_lst:
# process_slide(item, save_path='./result/pre_process/')