66 lines
2.6 KiB
Python
66 lines
2.6 KiB
Python
from PIL import Image
|
|
import glob
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
|
|
from utils import carregar_labelmap_completo, compute_roi_indices, resize_keep_width
|
|
|
|
# ----------------------------
|
|
# Dataset
|
|
# ----------------------------
|
|
class ROISegDataset(Dataset):
|
|
def __init__(self, root, out_dir, zona_inicio, faixa_atuacao, input_w=384, min_input_h=96, labelmap_path="labelmap.txt",
|
|
mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
|
|
|
|
self.out_dir = out_dir
|
|
self.img_paths = sorted(glob.glob(os.path.join(root, "images", "*")))
|
|
self.msk_paths = sorted(glob.glob(os.path.join(root, "masks", "*")))
|
|
assert len(self.img_paths) == len(self.msk_paths) and len(self.img_paths) > 0, \
|
|
f"Nenhuma imagem/máscara encontrada em {root}"
|
|
|
|
self.zona_inicio = zona_inicio
|
|
self.faixa_atuacao = faixa_atuacao
|
|
self.input_w = input_w
|
|
self.min_input_h = min_input_h
|
|
self.mean = np.array(mean, dtype=np.float32).reshape(1, 1, 3)
|
|
self.std = np.array(std, dtype=np.float32).reshape(1, 1, 3)
|
|
|
|
_, _, self.classes, self.ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
|
self.ignore_id = self.ignore_rgb[0]
|
|
|
|
def __len__(self):
|
|
return len(self.img_paths)
|
|
|
|
def __getitem__(self, idx):
|
|
img_rgb = np.array(Image.open(self.img_paths[idx]).convert("RGB"))
|
|
msk_grayscale = np.array(Image.open(self.msk_paths[idx]).convert("L"))
|
|
|
|
H, W = img_rgb.shape[:2]
|
|
y_fim, y_inicio = compute_roi_indices(H, self.zona_inicio, self.faixa_atuacao)
|
|
|
|
img_roi = img_rgb[y_fim:y_inicio, 0:W]
|
|
msk_roi = msk_grayscale[y_fim:y_inicio, 0:W]
|
|
|
|
img_in = resize_keep_width(img_roi, self.input_w, self.min_input_h)
|
|
msk_ids = resize_keep_width(msk_roi, self.input_w, self.min_input_h)
|
|
|
|
valores_validos = list(range(len(self.classes))) + [255]
|
|
msk_ids[np.isin(msk_ids, valores_validos, invert=True)] = self.ignore_id # converte inválidos em ignore
|
|
|
|
valores_unicos = np.unique(msk_ids)
|
|
if np.all(msk_ids == 255):
|
|
raise ValueError(f"Máscara {self.msk_paths[idx]} está só com valor de ignore (255)")
|
|
|
|
img_f = img_in.astype(np.float32) / 255.0
|
|
img_f = (img_f - self.mean) / self.std
|
|
img_chw = np.transpose(img_f, (2, 0, 1))
|
|
|
|
if idx == 0:
|
|
cv2.imwrite(os.path.join(self.out_dir, "debug_roi_input.png"), img_roi)
|
|
cv2.imwrite(os.path.join(self.out_dir, "debug_roi_mask.png"), msk_roi)
|
|
|
|
return torch.from_numpy(img_chw).float(), torch.from_numpy(msk_ids.astype(np.int64))
|