def __getitem__(self, idx):
imageid = self.image_ids[idx]
im1 = get_image(imageid, basepath=self.basepath, time='im1')
im2 = get_image(imageid, basepath=self.basepath, time='im2')
augmented = self.aug(image=im1, image1=im2)
return img_to_tensor(augmented['image']), img_to_tensor(
augmented['image1']), imageid
def __getitem__(self, idx):
name = self.names[idx]
img_path = os.path.join(self.data_path, self.image_type, name)
image = skimage.io.imread(img_path)
if image is None:
raise ValueError(img_path)
image_path = os.path.join(self.data_path, "train_mask_binned_mc_10", name[:-4].replace("RGB", "MS") + ".tif")
try:
raster_ds = gdal.Open(image_path, gdal.GA_ReadOnly)
speed_mask = raster_ds.ReadAsArray()
except:
# retry
raster_ds = gdal.Open(image_path, gdal.GA_ReadOnly)
speed_mask = raster_ds.ReadAsArray()
if speed_mask.shape[-1] > 100:
speed_mask = np.moveaxis(speed_mask, 0, -1)
junction_mask = cv2.imread(os.path.join(self.data_path, "junction_masks", name[:-4].replace("PS-RGB_", "").replace("PS-MS_", "") + ".png"), cv2.IMREAD_GRAYSCALE)
if junction_mask is None:
junction_mask = np.zeros((1300, 1300))
if speed_mask is None:
speed_mask = np.zeros((1300, 1300, 11), dtype=np.uint8)
sample = self.transforms(image=image, mask=np.concatenate([speed_mask, np.expand_dims(junction_mask, -1), np.expand_dims(junction_mask, -1)], axis=-1), img_name=name)
mask = sample["mask"]
city = np.zeros((len(cities), 1, 1))
city[get_city_idx(name), 0, 0] = 1
sample['city'] = city
sample['img_name'] = name
mask = np.moveaxis(mask, -1, 0) / 255.
sample['mask'] = torch.from_numpy(mask)
sample['image'] = img_to_tensor(sample["image"], self.normalize)
return sample
def __getitem__(self, idx):
img_file_name = self.file_names[idx]
image = load_image(img_file_name)
mask = load_mask(img_file_name, self.problem_type)
data = {"image": image, "mask": mask}
augmented = self.transform(**data)
image, mask = augmented["image"], augmented["mask"]
if self.mode == 'train':
if self.problem_type == 'binary':
return img_to_tensor(image), torch._C.from_numpy(
np.expand_dims(mask, 0)).float()
else:
return img_to_tensor(image), torch._C.from_numpy(mask).long()
else:
return img_to_tensor(image), str(img_file_name)
def __getitem__(self, idx):
imageid = self.image_ids[idx]
im1 = get_image(imageid, basepath=self.basepath, time='im1')
im2 = get_image(imageid, basepath=self.basepath, time='im2')
mask = cv2.imread(f'{self.basepath}/mask/{imageid}.png',
cv2.IMREAD_GRAYSCALE)
assert mask is not None
augmented = self.aug(image=im1, image1=im2, mask=mask)
mask_ = (augmented['mask'] > 0).astype(np.uint8)
mask_ = torch.from_numpy(np.expand_dims(mask_, 0)).float()
# label_ = torch.from_numpy(np.expand_dims(augmented['mask'], 0)).float()
return (img_to_tensor(augmented['image']),
img_to_tensor(augmented['image1']), mask_, imageid)
def __getitem__(self, idx):
image_path = self.image_df.iloc[idx].path
image = self._read_image(image_path, one_channel=False)
mask_path = self.mask_df.iloc[idx].path
if mask_path.exists():
mask = self._read_image(mask_path, one_channel=True)
else:
mask = np.zeros_like(image)[:, :, :1]
if self.transform:
augmented = self.transform(image=image, mask=mask)
image = augmented['image']
mask = augmented['mask']
image = img_to_tensor(image)
mask = img_to_tensor(mask)
return image, mask
def __getitem__(self, idx):
name = self.names[idx]
img_path = os.path.join(self.data_path, name)
image = skimage.io.imread(img_path)
if image is None:
raise ValueError(img_path)
sample = self.transforms(image=image)
sample['img_name'] = name
sample['image'] = img_to_tensor(sample["image"], self.normalize)
return sample
def pil_img_to_tensor(img: np.ndarray):
"""[Convert np.ndarray Image to torch.Tensor]
Args:
img (np.ndarray): [Image ]
Returns:
[torch.Tensor]: [returns processed image]
"""
img_tensor = T.Compose([
T.ToTensor(),
])
processed_img = img_to_tensor(img)
return processed_img
def predict_localization(image, config: ModelConfig):
conf = load_config(config.config_path)
model = models.__dict__[conf['network']](seg_classes=1,
backbone_arch=conf['encoder'])
checkpoint_path = os.path.join(weight_path, config.weight_path)
print("=> loading checkpoint '{}'".format(checkpoint_path))
checkpoint = torch.load(checkpoint_path, map_location="cpu")
model.load_state_dict({
k.replace("module.", ""): v
for k, v in checkpoint['state_dict'].items()
})
model.eval()
model = model.cpu()
print("predicting", config)
with torch.no_grad():
image = img_to_tensor(image, conf["input"]["normalize"]).cpu().numpy()
if "dpn" in config.weight_path:
image = np.pad(image, [(0, 0), (16, 16), (16, 16)], mode='reflect')
images = np.array([
image,
image[:, ::-1, :],
image[:, :, ::-1],
image[:, ::-1, ::-1],
])
images = torch.from_numpy(images).cpu().float()
prediction_masks = []
for i in range(4):
logits = model(images[i:i + 1])
preds = torch.sigmoid(logits).cpu().numpy()
pred = preds[0]
if i == 1:
pred = preds[0].copy()[:, ::-1, :]
if i == 2:
pred = preds[0].copy()[:, :, ::-1]
if i == 3:
pred = preds[0].copy()[:, ::-1, ::-1]
prediction_masks.append(pred)
preds = np.average(prediction_masks, axis=0)
if "dpn" in config.weight_path:
preds = preds[:, 16:-16, 16:-16]
return preds
def predict_damage(image, config: ModelConfig):
conf = load_config(config.config_path)
model = models.__dict__[conf['network']](seg_classes=5,
backbone_arch=conf['encoder'])
checkpoint_path = os.path.join(weight_path, config.weight_path)
print("=> loading checkpoint '{}'".format(checkpoint_path))
checkpoint = torch.load(checkpoint_path, map_location="cpu")
model.load_state_dict(
{k[7:]: v
for k, v in checkpoint['state_dict'].items()})
model.eval()
model = model.cpu()
print("predicting", config)
with torch.no_grad():
image = img_to_tensor(image, conf["input"]["normalize"]).cpu().numpy()
images = np.array([
image,
image[:, ::-1, :],
image[:, :, ::-1],
image[:, ::-1, ::-1],
])
images = torch.from_numpy(images).cpu().float()
prediction_masks = []
for i in range(4):
logits = model(images[i:i + 1])
preds = torch.softmax(logits, dim=1).cpu().numpy()
pred = preds[0]
if i == 1:
pred = preds[0].copy()[:, ::-1, :]
if i == 2:
pred = preds[0].copy()[:, :, ::-1]
if i == 3:
pred = preds[0].copy()[:, ::-1, ::-1]
prediction_masks.append(pred)
preds = np.average(prediction_masks, axis=0)
return preds
w = math.ceil(w / factor) * factor // 2
#print(h, w)
mask = np.zeros(shape=(h, w))
image = cv2.resize(image, (w, h))
image_ori = image
# image=illum(image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
aug = Compose([
# PadIfNeeded(min_height=h, min_width=w, border_mode=cv2.BORDER_CONSTANT, value=0, p=1.0), # padding到2的5次方的倍数
Normalize(p=1) # 归一化
])
augmented = aug(image=image, mask=mask)
image = augmented['image']
image = img_to_tensor(image).unsqueeze(0).to(
DEVICE
) # torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float().to(DEVICE) # 图像转为tensor格式
output = model(image) # 预测
seg_mask = (output[0].data.cpu().numpy().argmax(axis=0)).astype(
np.uint8)
t2 = time.time()
print("time:", (t2 - t1))
full_mask = np.zeros((h, w, 3))
for mask_label, sub_color in enumerate(class_color):
full_mask[seg_mask == mask_label, 0] = sub_color[2]
full_mask[seg_mask == mask_label, 1] = sub_color[1]
full_mask[seg_mask == mask_label, 2] = sub_color[0]
# print(full_mask.max())
# import matplotlib.pyplot as plt
# plt.imshow(full_mask)
def start_divide(self):
class_color = [[0, 0, 0], [0, 255, 0], [0, 255, 255], [125, 255, 12]]
t1 = time.time()
# ********
seg = "binary" # 多类别分割，夹子为一类，铰链为一类，柄为一类
if seg == "parts":
model = parts_seg(in_channels=3, out_channels=4)
# model_path = r"model_2_TDSNet.pt" # 多类别分割模型地址
elif seg == "binary": # 二元分割，背景为一类手术器械为一类
model = binary_seg(in_channels=3, out_channels=2)
# model_path = r"model_0_TDSNet.pt" # 二元分割模型地址
else:
model = r""
model_path = r""
model_path = self.model_path
state = torch.load(str(model_path), map_location='cpu')
state = {
key.replace('module.', ''): value
for key, value in state['model'].items()
}
model.load_state_dict(state)
model.to(device)
model.eval()
# **********
types = self.ui.type_combo.currentText()
path = self.ui.data_path_text.toPlainText() # 要分割的图像或者视频或者nii所在的文件夹地址
path_save = self.ui.result_path_text.toPlainText() # 要保存图像的地址
names = file_name(path) # path目录下所有文件名字
factor = 2**5
# ********
if types == "image":
for i, name in enumerate(names):
t1 = time.time()
path_single = os.path.join(path, name)
image = cv2.imdecode(np.fromfile(path_single, dtype=np.uint8),
-1)
h, w, channel = image.shape
h = math.ceil(
h / factor
) * factor // 2 # 向上取整，由于模型需要下采样5次图像会变成原来的2的5次方分之一，需要输入图像是2的5次方的倍数
w = math.ceil(w / factor) * factor // 2
# print(h, w)
mask = np.zeros(shape=(h, w))
image = cv2.resize(image, (w, h))
image_ori = image
# image=illum(image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
aug = Compose([
# PadIfNeeded(min_height=h, min_width=w, border_mode=cv2.BORDER_CONSTANT, value=0, p=1.0), # padding到2的5次方的倍数
Normalize(p=1) # 归一化
])
augmented = aug(image=image, mask=mask)
image = augmented['image']
image = img_to_tensor(image).unsqueeze(0).to(
device
) # torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float().to(DEVICE) # 图像转为tensor格式
output = model(image) # 预测
seg_mask = (output[0].data.cpu().numpy().argmax(
axis=0)).astype(np.uint8)
t2 = time.time()
print("time:", (t2 - t1))
full_mask = np.zeros((h, w, 3))
for mask_label, sub_color in enumerate(class_color):
full_mask[seg_mask == mask_label, 0] = sub_color[2]
full_mask[seg_mask == mask_label, 1] = sub_color[1]
full_mask[seg_mask == mask_label, 2] = sub_color[0]
# print(full_mask.max())
# import matplotlib.pyplot as plt
# plt.imshow(full_mask)
# plt.show()
seg = mask_overlay(image_ori, (full_mask > 0)).astype(np.uint8)
# seg = mask_overlay((full_mask > 0)).astype(np.uint8)
# cv2.imshow("seg",seg)
# cv2.waitKey(0)
cv2.imwrite(path_save + "/" + str(i).zfill(4) + ".png", seg)
save_video(path_save, path_save, h, w,
self.aviname) # 可以播放保存的视频展示分割结果
elif types == "video":
# 对视频进行分割
splitFrames(path, path_save)
# 调用图像分割的方法分割从视频中保存的图像
elif types == "nii":
# 对nii3维医学图像进行分割
import nibabel as nib
img = np.array(nib.load(path).get_data())
pass
from dataset import load_image
import torch
from generate_masks import get_model, img_transform
from functions import jaccard
import numpy as np
import matplotlib.pyplot as plt
from albumentations.pytorch.transforms import img_to_tensor
arrt = np.random.randn(500, 500)
tensor = img_to_tensor(arrt)
mean=(0.485, 0.456, 0.406)
std=(0.229, 0.224, 0.225)
#new_img = img - mean /std
#old img =
model_path = 'pretrained/unet16_instruments_20/model_1.pt'
model = get_model(model_path, model_type='UNet16', problem_type='instruments')
img_l = load_image('dataset/instrument_dataset_1/left_frames/frame000.png')
input_img_l = torch.unsqueeze(img_to_tensor(img_transform(p=1)(image=img_l)['image']), dim=0)
mask_l = model(input_img_l)
img_r = load_image('dataset/instrument_dataset_1/right_frames/frame000.png')
input_img_r = torch.unsqueeze(img_to_tensor(img_transform(p=1)(image=img_r)['image']), dim=0)
mask_r = model(input_img_r)
mask = jaccard(mask_l, mask_r)
im_seg = mask.data[0].cpu().numpy()[0]
