246 lines
11 KiB
Python
246 lines
11 KiB
Python
import os
|
|
import re
|
|
import numpy as np
|
|
import nrrd
|
|
import pydicom
|
|
from pydicom.uid import ExplicitVRLittleEndian, generate_uid
|
|
from pathlib import Path
|
|
import matplotlib
|
|
matplotlib.use('Agg')
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.patches import Rectangle
|
|
from matplotlib.colors import ListedColormap
|
|
from scipy import ndimage
|
|
import cv2
|
|
|
|
# === НАСТРОЙКИ ===
|
|
BASE_DIR = "data/1_2236040_zno_cyst_AF"
|
|
IMG_NRRD = "501 arterial ax 1mm iDose (5).nrrd"
|
|
SEG_NRRD = "Segmentation.seg.nrrd"
|
|
CT_CMAP = 'gray'
|
|
|
|
# Параметр аппроксимации полигона (0.01 = 1% от длины контура)
|
|
POLY_EPSILON = 0.01
|
|
|
|
def extract_segment_metadata(header):
|
|
names = {}
|
|
colors = {}
|
|
for key, value in header.items():
|
|
if key.startswith('Segment') and key.endswith('_LabelValue'):
|
|
match = re.search(r'Segment(\d+)_LabelValue', key)
|
|
if match:
|
|
seg_idx = int(match.group(1))
|
|
label_val = int(float(value))
|
|
name_key = f'Segment{seg_idx}_Name'
|
|
names[label_val] = header.get(name_key, f'Class_{label_val}').strip()
|
|
color_key = f'Segment{seg_idx}_Color'
|
|
color_str = header.get(color_key, '')
|
|
if color_str:
|
|
parts = color_str.strip().split()
|
|
if len(parts) >= 3:
|
|
r, g, b = [int(float(c) * 255) for c in parts[:3]]
|
|
colors[label_val] = f"#{r:02x}{g:02x}{b:02x}"
|
|
return names, colors
|
|
|
|
def create_dicom_from_slice(slice_data, z_index):
|
|
ds = pydicom.Dataset()
|
|
ds.file_meta = pydicom.Dataset()
|
|
ds.file_meta.TransferSyntaxUID = ExplicitVRLittleEndian
|
|
ds.file_meta.MediaStorageSOPClassUID = '1.2.840.10008.5.1.4.1.1.2'
|
|
ds.file_meta.MediaStorageSOPInstanceUID = generate_uid()
|
|
ds.file_meta.ImplementationClassUID = generate_uid()
|
|
ds.SOPClassUID = ds.file_meta.MediaStorageSOPClassUID
|
|
ds.SOPInstanceUID = ds.file_meta.MediaStorageSOPInstanceUID
|
|
ds.PatientName = "ANON^PATIENT"
|
|
ds.PatientID = "000000"
|
|
ds.StudyDate = "20260602"
|
|
ds.Modality = "CT"
|
|
ds.InstanceNumber = str(z_index)
|
|
ds.Rows, ds.Columns = slice_data.shape
|
|
ds.PixelSpacing = [1.0, 1.0]
|
|
ds.SliceThickness = 1.0
|
|
ds.BitsAllocated = 16; ds.BitsStored = 16; ds.HighBit = 15; ds.PixelRepresentation = 1
|
|
ds.ImagePositionPatient = [0.0, 0.0, float(z_index)]
|
|
ds.ImageOrientationPatient = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]
|
|
ds.SamplesPerPixel = 1; ds.PhotometricInterpretation = "MONOCHROME2"
|
|
ds.PixelData = slice_data.tobytes()
|
|
return ds
|
|
|
|
def extract_polygon(mask_uint8):
|
|
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
if not contours:
|
|
return None
|
|
largest_contour = max(contours, key=cv2.contourArea)
|
|
epsilon = POLY_EPSILON * cv2.arcLength(largest_contour, True)
|
|
approx = cv2.approxPolyDP(largest_contour, epsilon, True)
|
|
if len(approx) < 3:
|
|
return None
|
|
points = [(int(pt[0][0]), int(pt[0][1])) for pt in approx]
|
|
return points
|
|
|
|
def save_png_safe(path: Path, image: np.ndarray):
|
|
"""
|
|
Надёжное сохранение PNG через cv2.imencode.
|
|
Обходит ограничение cv2.imwrite на Unicode-пути в Windows.
|
|
"""
|
|
success, buffer = cv2.imencode('.png', image)
|
|
if not success:
|
|
raise RuntimeError(f"cv2.imencode failed for {path}")
|
|
path.write_bytes(buffer.tobytes())
|
|
|
|
def main():
|
|
base_path = Path(BASE_DIR)
|
|
if not base_path.exists():
|
|
print(f"❌ Ошибка: Директория '{BASE_DIR}' не найдена.")
|
|
return
|
|
|
|
print(f"📂 Рабочая директория: {base_path.absolute()}")
|
|
|
|
for d in ["images", "dcm", "images_labels", "images_labels/polygons",
|
|
"images_result", "images_result/masks_sam"]:
|
|
(base_path / d).mkdir(exist_ok=True)
|
|
|
|
print("📥 Загрузка NRRD...")
|
|
img_data, _ = nrrd.read(base_path / IMG_NRRD)
|
|
seg_data, seg_header = nrrd.read(base_path / SEG_NRRD)
|
|
|
|
if img_data.shape != seg_data.shape:
|
|
raise ValueError(f"Размеры не совпадают: {img_data.shape} vs {seg_data.shape}")
|
|
|
|
SEGMENT_NAMES, SEGMENT_COLORS = extract_segment_metadata(seg_header)
|
|
if not SEGMENT_NAMES:
|
|
print("️ Метаданные не найдены. Используются заглушки.")
|
|
unique = np.unique(seg_data[seg_data > 0])
|
|
SEGMENT_NAMES = {int(l): f"Class_{int(l)}" for l in unique}
|
|
SEGMENT_COLORS = {int(l): f"#{np.random.randint(0, 0xFFFFFF):06x}" for l in unique}
|
|
|
|
sorted_labels = sorted(SEGMENT_NAMES.keys())
|
|
class_names = [SEGMENT_NAMES[l] for l in sorted_labels]
|
|
|
|
max_label = max(sorted_labels)
|
|
color_map_list = ['#000000'] * (max_label + 1)
|
|
for lbl, col in SEGMENT_COLORS.items():
|
|
if lbl <= max_label:
|
|
color_map_list[lbl] = col
|
|
custom_cmap = ListedColormap(color_map_list)
|
|
|
|
print(f"✅ Найдено {len(class_names)} классов: {class_names}")
|
|
(base_path / "predefined_classes.txt").write_text("\n".join(class_names))
|
|
|
|
W, H, D = img_data.shape
|
|
count = 0
|
|
mask_count = 0 # <-- счётчик масок для диагностики
|
|
|
|
for z in range(D):
|
|
img_slice = img_data[:, :, z].T
|
|
seg_slice = seg_data[:, :, z].T
|
|
|
|
create_dicom_from_slice(img_slice.astype(np.int16), z).save_as(base_path / "dcm" / f"slice_{z:04d}.dcm")
|
|
|
|
p_min, p_max = np.percentile(img_slice, [0.5, 99.5])
|
|
img_uint8 = (np.clip(img_slice, p_min, p_max) - p_min) / (p_max - p_min + 1e-8) * 255
|
|
img_uint8 = img_uint8.astype(np.uint8)
|
|
|
|
slice_labels = np.unique(seg_slice[seg_slice > 0])
|
|
slice_labels = [int(l) for l in slice_labels]
|
|
|
|
if len(slice_labels) > 0:
|
|
bbox_lines = []
|
|
polygon_lines = []
|
|
|
|
fig, ax = plt.subplots(figsize=(img_slice.shape[1]/100, img_slice.shape[0]/100))
|
|
ax.imshow(img_uint8, cmap=CT_CMAP)
|
|
ax.imshow(seg_slice, cmap=custom_cmap, alpha=0.4)
|
|
|
|
H_curr, W_curr = img_slice.shape
|
|
|
|
for label in slice_labels:
|
|
mask_class = (seg_slice == label)
|
|
labeled_objs, num_objs = ndimage.label(mask_class, structure=np.ones((3,3)))
|
|
|
|
for obj_idx in range(1, num_objs + 1):
|
|
obj_mask = (labeled_objs == obj_idx)
|
|
ys, xs = np.where(obj_mask)
|
|
if len(ys) == 0: continue
|
|
|
|
x_min, y_min, x_max, y_max = xs.min(), ys.min(), xs.max(), ys.max()
|
|
|
|
# --- BBOX (YOLO) ---
|
|
xc = (x_min + x_max) / 2.0
|
|
yc = (y_min + y_max) / 2.0
|
|
bw = x_max - x_min
|
|
bh = y_max - y_min
|
|
yolo_class_id = sorted_labels.index(label)
|
|
bbox_lines.append(f"{yolo_class_id} {xc/W_curr:.6f} {yc/H_curr:.6f} {bw/W_curr:.6f} {bh/H_curr:.6f}")
|
|
|
|
# --- POLYGON (YOLO Segmentation) ---
|
|
obj_mask_uint8 = (obj_mask * 255).astype(np.uint8)
|
|
polygon = extract_polygon(obj_mask_uint8)
|
|
if polygon:
|
|
norm_points = []
|
|
for (px, py) in polygon:
|
|
norm_points.append(f"{px/W_curr:.6f}")
|
|
norm_points.append(f"{py/H_curr:.6f}")
|
|
polygon_str = " ".join(norm_points)
|
|
polygon_lines.append(f"{yolo_class_id} {polygon_str}")
|
|
|
|
# --- SAM Mask (ИСПРАВЛЕНО: надёжное сохранение) ---
|
|
single_mask = np.zeros((H_curr, W_curr), dtype=np.uint8)
|
|
single_mask[obj_mask] = 255
|
|
mask_filename = f"slice_{z:04d}_class_{yolo_class_id}_obj_{(obj_idx-1):04d}.png"
|
|
mask_path = base_path / "images_result" / "masks_sam" / mask_filename
|
|
|
|
try:
|
|
save_png_safe(mask_path, single_mask)
|
|
mask_count += 1
|
|
except Exception as e:
|
|
print(f"⚠️ Ошибка сохранения маски {mask_path}: {e}")
|
|
|
|
# --- Отрисовка на превью ---
|
|
color_hex = color_map_list[label]
|
|
rect = Rectangle((x_min, y_min), bw, bh, linewidth=2,
|
|
edgecolor=color_hex, facecolor='none', linestyle='-', alpha=0.7)
|
|
ax.add_patch(rect)
|
|
if polygon:
|
|
poly_pts = np.array(polygon)
|
|
closed_pts = np.vstack([poly_pts, poly_pts[0:1]])
|
|
ax.plot(closed_pts[:, 0], closed_pts[:, 1],
|
|
color=color_hex, linewidth=2.5, linestyle='--', alpha=0.9)
|
|
|
|
(base_path / "images_labels" / f"slice_{z:04d}.txt").write_text("\n".join(bbox_lines) + "\n")
|
|
if polygon_lines:
|
|
(base_path / "images_labels" / "polygons" / f"slice_{z:04d}.txt").write_text("\n".join(polygon_lines) + "\n")
|
|
|
|
plt.imsave(base_path / "images" / f"slice_{z:04d}.png", img_uint8, cmap=CT_CMAP, format='png')
|
|
|
|
ax.set_title(f"Slice {z} | BBoxes: {len(bbox_lines)} | Polys: {len(polygon_lines)}")
|
|
ax.axis('off')
|
|
plt.tight_layout(pad=0)
|
|
plt.savefig(base_path / "images_result" / f"slice_{z:04d}_preview.png", dpi=150, bbox_inches='tight')
|
|
plt.close(fig)
|
|
count += 1
|
|
|
|
(base_path / "data.yaml").write_text(f"""path: {base_path.absolute()}
|
|
train: images
|
|
val: images
|
|
test: images
|
|
nc: {len(class_names)}
|
|
names: {class_names}
|
|
""")
|
|
|
|
print(f"\n✅ Готово! Обработано срезов: {count}")
|
|
print(f"💾 Создано бинарных масок: {mask_count}")
|
|
print(f" 📁 images/ → PNG для обучения")
|
|
print(f" 📁 images_labels/ → YOLO bbox (.txt)")
|
|
print(f" └── polygons/ → YOLO polygon segmentation (.txt)")
|
|
print(f" └── predefined_classes.txt")
|
|
print(f" 📁 images_result/ → Цветные превью")
|
|
print(f" └── masks_sam/ → Бинарные маски 0/255 ({mask_count} шт.)")
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except Exception as e:
|
|
print(f"❌ Ошибка: {e}")
|
|
import traceback
|
|
traceback.print_exc() |