iunknown/ai-med-edu
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: project structure, scripts and documentation

Browse Source
This commit is contained in:
iunknown
2026-08-22 13:17:29 +03:00
commit 5042305df0
54 changed files with 14166 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/Тема 2.2/nrrd_to_yolo.py
+246
View File
@@ -0,0 +1,246 @@
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()