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()