seg.py

from dino.predict import predict_img,mask_to_image
import torch
from PIL import Image
import numpy as np
import cv2

# 송장부분을 crop 하기 위한 코드 
# dinov2 모델을 segmentation 모델로 활용하였고 이를 train 하여서 pretrained weight를 weight 폴더에 저장하고 아래와 같이 load합니다.

def segmentation(img):
    #device=torch.device("cuda")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net=torch.load("weights/dinov2_model.pth")
    
    mask_values=[[0, 0, 0], [255, 255, 255]]

    mask=predict_img(net,img,device,scale_factor=1,out_threshold=0.5)
    result = mask_to_image(mask, mask_values)
    result=np.array(result)
    
    
    return result 


# def calculate_white_area_percentage(mask):
#     total_pixels = mask.size
#     print(total_pixels)
    
#     white_pixels = np.sum(np.all(mask >100, axis=-1))
    
#     white_area_percentage = (white_pixels / total_pixels) * 100
    
#     return white_area_percentage*3

    
# 위 segmentation 을 통해서 crop 된 부분이 이미지내에서 몇프로 차지하는지 계산합니다. 
# 아래 코드는 하얀색 픽셀이 연속적으로 이어져서 만들어진 덩어리가 전체에서 몇프로 차지하는지 계산합니다. 
# 아래 코드는 덩어리(송장으로 추정) 들이 2개 이상이어도 적용할수 있습니다. 
def mask_percentage(mask_path):

    image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

    ret, threshold = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY)

    contours, hierarchy = cv2.findContours(threshold, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    total_area = image.shape[0] * image.shape[1]  
    contours_list=contours
    
    contour_areas = [cv2.contourArea(contour) for contour in contours]
    

    percentages = [(area / total_area) * 100 for area in contour_areas]
    percentage_list=[]
    for i, percentage in enumerate(percentages):
        percentage_list.append(percentage)
    return contours_list,percentage_list