def preprocess_frames(img, out_path, idx):
if (type(img) != np.ndarray and type(img) != np.memmap) or img is None:
return None
img_enh = utl.enhance_contrast_image(img, clip_limit=3.5, tile_size=12)
mask, circle = utl.get_retina_mask(img_enh)
if circle[2] == 0:
return None
img = cv2.bitwise_and(img, mask)
img = utl.crop_to_circle(img, circle)
if img is not None and img.size != 0:
cv2.imwrite(join(out_path, SUBFOLDER_PROCESSED, f'{idx}.jpg'), img)
return img
def preprocess_images(image_path, output_path, min_radius=300):
img = cv2.imread(image_path)
if (type(img) != np.ndarray and type(img) != np.memmap) or img is None:
return
img_enh = utl.enhance_contrast_image(img, clip_limit=3.5, tile_size=12)
mask, circle = utl.get_retina_mask(img_enh, min_radius=min_radius)
if circle[2] == 0:
return
img = cv2.bitwise_and(img, mask)
img = utl.crop_to_circle(img, circle)
if img is not None and img.size != 0:
filepath = join(output_path, os.path.splitext(os.path.basename(image_path))[0]) + '.jpg'
cv2.imwrite(filepath, img)
def run(input_path, model_path, process):
# Loading image
img = cv2.imread(input_path)
if (type(img) != np.ndarray and type(img) != np.memmap) or img is None:
print('Invalid input image: ', input_path)
return
if process:
img_enh = enhance_contrast_image(img, clip_limit=3.5, tile_size=12)
mask, circle = get_retina_mask(img_enh)
if circle[2] == 0:
print('Could not detect retinoscope lens.')
return
img = cv2.bitwise_and(img, mask)
img = crop_to_circle(img, circle)
print('Loaded image successfully...')
# Necessary image augmentations
aug_pipeline = alb.Compose([
alb.Resize(425, 425, always_apply=True, p=1.0),
alb.CenterCrop(399, 399, always_apply=True, p=1.0),
alb.Normalize(always_apply=True, p=1.0),
ToTensorV2(always_apply=True, p=1.0)
],
p=1.0)
img_tensor = aug_pipeline(image=img)['image']
# Loading model
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
model = models.alexnet()
model.classifier[-1] = nn.Linear(model.classifier[-1].in_features, 2)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
print('Loaded model successfully...')
#Prediction
id_to_class = {0: 'no_retinopathy', 1: 'retinopathy'}
outputs = model(img_tensor.unsqueeze(0))
_, prediction = torch.max(outputs, 1)
prediction = prediction[0].item()
print()
print(
f'Prediction for image {os.path.basename(input_path)} with the model {os.path.basename(model_path)}:'
)
print(f'------ {id_to_class[prediction]} ------')