atom-predict/egnn_v2/pre_lb.py

62 lines
1.5 KiB
Python

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 = {
'Norm':1,
'SV':2,
'LineSV':3,
}
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
# Image.fromarray(img).save(
# os.path.join(save_path, '{}.png'.format(base_name))
# )
Image.fromarray(mask).save(
os.path.join(save_path, 'lbl', '{}.png'.format(base_name))
)
img_lst = glob.glob('/home/gao/mouclear/cc/final_todo/SV/end-to-end_test/*.jpg')
for item in img_lst:
process_slide(item, save_path='/home/gao/mouclear/cc/final_todo/SV/end-to-end_test/')