ObjectDetection/functions.py

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