Использование pydensecrf для масок, возвращаемых AI ⇐ Python

[Anonymous]

Обученный искусственный интеллект обрабатывает изображение с помощью модуля YOLO. он возвращает объект, который я вставляю в функцию ниже. структура объектов – это

• результаты
  
  маски
  
  < li>xy — полигоны
• данные — бинарные маски

[*]блоки

• список диктов
  
  по ключевому индексу cls 0, содержащий идентификатор типа маски

Моя цель — обработать маски с помощью crf. код ниже дает выходные данные, но я не могу определить разницу между исходной маской и обработанными CRF.
мой вопрос: правильный ли мой код ниже? и что именно должен делать crf? (потому что я понимаю, что это лучше соответствует маскам)
import cv2
import numpy as np
import pydensecrf.densecrf as dcrf
import torch
from pydensecrf.utils import unary_from_softmax, create_pairwise_gaussian, create_pairwise_bilateral
from collections import defaultdict

def process_image_with_crf(image, ai_results, sdims_gaussian=(3, 3), sdims_bilateral=(50, 50), schan_bilateral=(20, 20, 20), compat_gaussian=3, compat_bilateral=10):
H, W, _ = image.shape

# Initialize result dictionary
results_by_class = defaultdict(lambda: {"masks": [], "polygons": []})

# Iterate through each mask and process individually
for binary_mask, polygon, box in zip(ai_results.masks.data.cpu().numpy(), ai_results.masks.xy, ai_results.boxes):
class_id = int(box.cls[0])

# Prepare softmax for CRF: two channels (foreground & background)
softmax_probs = np.zeros((H, W, 2), dtype=np.float32)
softmax_probs[:, :, 1] = binary_mask # Foreground (object)
softmax_probs[:, :, 0] = 1 - binary_mask # Background

# Normalize the softmax probabilities
softmax_probs /= np.sum(softmax_probs, axis=-1, keepdims=True)

# Reshape for unary potential (CRF expects shape (2, H * W))
softmax_probs = softmax_probs.transpose(2, 0, 1).reshape(2, -1)

# CRF processing
unary = unary_from_softmax(softmax_probs)
unary = np.ascontiguousarray(unary)
d = dcrf.DenseCRF2D(W, H, 2)
d.setUnaryEnergy(unary)

# Add Gaussian pairwise potential (spatial smoothness)
pairwise_gaussian = create_pairwise_gaussian(sdims=sdims_gaussian, shape=(H, W))
d.addPairwiseEnergy(pairwise_gaussian, compat=compat_gaussian)

# Add Bilateral pairwise potential (appearance consistency)
pairwise_bilateral = create_pairwise_bilateral(sdims=sdims_bilateral, schan=schan_bilateral, img=image, chdim=2)
d.addPairwiseEnergy(pairwise_bilateral, compat=compat_bilateral)

# Perform CRF inference
refined_mask = np.argmax(d.inference(10), axis=0).reshape((H, W))
refined_binary_mask = (refined_mask == 1).astype(np.uint8) * 255 # Convert to black-and-white mask

# Store results
results_by_class[class_id]["masks"].append(refined_binary_mask) # Refined mask
results_by_class[class_id]["polygons"].append(polygon) # Original polygon

return results_by_class


Подробнее здесь: https://stackoverflow.com/questions/793 ... rned-by-ai