File size: 4,320 Bytes
bfabfea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import torch
import cv2
import numpy as np
def cropBlackBackground(img):
"""Removes black background of labbel image and returns portion which contains label only.
It will need gray image.
Args:
img (ndarray): Numpy array representation of image
Returns:
img (ndarray): Numpy array representation of cropped image
"""
try:
_, binary = cv2.threshold(img, 1, 255, cv2.THRESH_BINARY)
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
max_area_indx = 0
max_area = 0
for indx, contour in enumerate(contours):
area = cv2.contourArea(contour)
if area > max_area:
max_area_indx = indx
x, y, w, h = cv2.boundingRect(contours[max_area_indx])
img = img[y:y+h, x:x+w]
return img
except Exception as e:
print(e)
return None
def enhanceImage(img, block_size: int = 19, constant: int = 5, adaptive_thresold_type = "GAUSSIAN", need_to_sharp: bool = True):
"""Enhance image by appling adaptive thresolding and filter2D
Args:
img (ndarray): Numpy array representation of image
block_size (int, optional): Block size for adaptive thresolding. Defaults to 25.
constant (int, optional): Constant for adaptive thresolding. Defaults to 10.
adaptive_thresold_type (str, optional): "GAUSSIAN" or "MEAN. Defaults to "GAUSSIAN".
need_to_sharp (bool, optional): Defaults to True.
Returns:
img (ndarray): Numpy array representation of enhanced image
"""
try:
if block_size < 2:
block_size = 2
block_size = block_size + 1 if block_size%2 == 0 else block_size
final_img = img
if adaptive_thresold_type == "MEAN":
final_img = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,block_size,constant)
else:
final_img = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,block_size,constant)
if need_to_sharp:
kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
final_img = cv2.filter2D(final_img, -1, kernel)
return final_img
except Exception as e:
print(e)
return None
def generateMask(res, original_img):
"""_summary_
Args:
res: Resulf of yolo image segmentation for single mask
original_img (ndarray): Numpy array representation of original image
Returns:
tupple (ndarray, ndarray): (crop_img, mask)
"""
try:
height,width = original_img.shape
masks = res.masks.data
boxes = res.boxes.data
#get index of box which has maximum confidence
max_conf_index = res.boxes.conf.argmax()
# extract classes
clss = boxes[:, 5]
# get indices of ress where class is 0
label_indices = torch.where(clss == 0)
# use these indices to extract the relevant masks
label_masks = masks[label_indices]
# get maximum confidence label's mask
max_conf_label_mask = torch.empty(size=(1,label_masks[max_conf_index].shape[0],label_masks[max_conf_index].shape[1]))
max_conf_label_mask[0]= label_masks[max_conf_index]
# scale for visualizing ress
label_mask = torch.any(max_conf_label_mask, dim=0).int() * 255
#final mask
final_mask = label_mask.cpu().numpy()
height_mask,width_mask =final_mask.shape
fy = height/height_mask
fx = width/width_mask
final_mask = cv2.resize(final_mask,(0,0),fx =fx ,fy = fy,interpolation=cv2.INTER_NEAREST)
original_img = original_img.astype(np.uint8)
final_mask = final_mask.astype(np.uint8)
# Expand boundries
kernel = np.ones((40,40), np.uint8)
expanded_mask = cv2.dilate(final_mask, kernel)
#crop_img
crop_img = cv2.bitwise_and(original_img,original_img,mask=expanded_mask)
return crop_img, expanded_mask
except Exception as e:
print(e)
return None, None
|