{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":113558,"databundleVersionId":14174843,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Convert Binary Mask Dataset to YOLO Segmentation Format\n\nFor Recod.ai/LUC - Scientific Image Forgery Detection Competition\n\nhttps://docs.ultralytics.com/datasets/segment/#ultralytics-yolo-format\n\n### Dataset Structure Expected:\n- Images folder: Contains original images\n- Masks folder: Contains binary masks (white=forgery, black=background)\n\n### YOLO Segmentation Format:\n- Each image has a corresponding .txt file\n- Format: <class_id> <x1> <y1> <x2> <y2> ... <xn> <yn>\n- Coordinates are normalized (0-1) relative to image dimensions","metadata":{}},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\n\nINPUT_DIR = \"/kaggle/input/recodai-luc-scientific-image-forgery-detection\"\n\n# Path to the training forged images and masks\ntrain_forged_images_dir = os.path.join(INPUT_DIR, 'train_images', 'forged')\n# Corrected path to masks directory\ntrain_masks_dir = os.path.join(INPUT_DIR, 'train_masks')\n\n# Get a list of sample forged images\nsample_image_name = os.listdir(train_forged_images_dir)[0]\nsample_image_path = os.path.join(train_forged_images_dir, sample_image_name)\n\n# Assuming the mask has the same name as the image\nsample_mask_name = sample_image_name.replace('.png', '.npy') # Masks are .npy files\nsample_mask_path = os.path.join(train_masks_dir, sample_mask_name)\n\ndef load_mask(mask_path: str) -> np.ndarray:\n    \"\"\"Load and binarize mask from .npy file.\"\"\"\n    mask_raw = np.load(mask_path)\n    # Binarize the 2D mask: 1 if value > 0, 0 otherwise.\n    mask = np.sum(mask_raw, axis=0)\n    return mask > 0\n\n\n# Load the sample image and mask\nsample_image = cv2.imread(sample_image_path)\nsample_mask = load_mask(sample_mask_path)\n\nprint(f\"Loaded sample image: {sample_image_name}\")\nprint(f\"Loaded sample mask: {sample_mask_name}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-29T08:26:56.092637Z","iopub.execute_input":"2025-10-29T08:26:56.092862Z","iopub.status.idle":"2025-10-29T08:26:56.498242Z","shell.execute_reply.started":"2025-10-29T08:26:56.092840Z","shell.execute_reply":"2025-10-29T08:26:56.497268Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\n# Display the image and mask side by side\nfig, axes = plt.subplots(1, 2, figsize=(12, 6))\n\naxes[0].imshow(cv2.cvtColor(sample_image, cv2.COLOR_BGR2RGB)) # Convert BGR to RGB for matplotlib\naxes[0].set_title(\"Sample Image\")\naxes[0].axis('off')\n\naxes[1].imshow(sample_mask) # Multiply by 255 is done in the line above\naxes[1].set_title(\"Sample Mask\")\naxes[1].axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-29T08:26:56.499558Z","iopub.execute_input":"2025-10-29T08:26:56.499846Z","iopub.status.idle":"2025-10-29T08:26:57.005879Z","shell.execute_reply.started":"2025-10-29T08:26:56.499822Z","shell.execute_reply":"2025-10-29T08:26:57.004825Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def mask_to_yolo_segmentation(mask: np.ndarray, tolerance: float = 0.005) -> list:\n    \"\"\"\n    Converts a binary mask to YOLO segmentation format with contour approximation.\n\n    Args:\n        mask: A 2D numpy array representing the binary mask (1 for foreground, 0 for background).\n        tolerance: The approximation accuracy. A smaller value gives a more precise approximation\n                   with more vertices.\n\n    Returns:\n        A string in YOLO segmentation format: \"<class_id> x1 y1 x2 y2 ...\".\n        Returns an empty string if no foreground pixels are found.\n    \"\"\"\n    # Convert the mask to uint8\n    mask_uint8 = mask.astype(np.uint8)\n    # Assuming mask is the same size as the image\n    height, width = mask.shape\n\n    # Find contours in the mask\n    contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n    yolo_contours = []\n\n    for contour in contours:\n        # Approximate the contour\n        epsilon = tolerance * cv2.arcLength(contour, True)\n        approx_contour = cv2.approxPolyDP(contour, epsilon, True)\n\n        # Reshape approximated contour for YOLO format\n        segmentation = approx_contour.flatten().tolist()\n\n        # Normalize the coordinates\n        normalized_segmentation = [\n            segmentation[i] / (width if i % 2 == 0 else height) for i in range(len(segmentation))\n        ]\n        yolo_contours.append(normalized_segmentation)\n\n        # Format the string\n        #class_id = 1\n        #segmentation_string = \" \".join([f\"{coord:.6f}\" for coord in normalized_segmentation])\n        #yolo_format_string += f\"{class_id} {segmentation_string}\\n\"\n\n    return yolo_contours\n\n\ndef format_yolo_segmentation(yolo_contours: list, class_id: int = 0) -> str:\n    \"\"\"\n    Formats YOLO segmentation contours into a string.\n\n    Args:\n        yolo_contours: A list of YOLO segmentation contours.\n        class_id: The class ID to use in the format string.\n\n    Returns:\n        A string in YOLO segmentation format: \"<class_id> x1 y1 x2 y2 ...\".\n    \"\"\"\n    yolo_format_string = \"\"\n    for contour in yolo_contours:\n        segmentation_string = \" \".join([f\"{coord:.6f}\" for coord in contour])\n        yolo_format_string += f\"{class_id} {segmentation_string}\\n\"\n    return yolo_format_string.strip() # Remove trailing newline\n\n\n# Convert the sample mask to YOLO segmentation format\nsample_yolo_annotation = mask_to_yolo_segmentation(sample_mask)\n\nprint(\"Sample YOLO Annotation:\")\nprint(format_yolo_segmentation(sample_yolo_annotation))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-29T08:26:57.007083Z","iopub.execute_input":"2025-10-29T08:26:57.007397Z","iopub.status.idle":"2025-10-29T08:26:57.027554Z","shell.execute_reply.started":"2025-10-29T08:26:57.007370Z","shell.execute_reply":"2025-10-29T08:26:57.026703Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def draw_yolo_segmentation(image: np.ndarray, yolo_annotation: list, color=(0, 255, 0), thickness=2) -> np.ndarray:\n    \"\"\"\n    Draws YOLO segmentation contours on an image.\n\n    Args:\n        image: The input image (NumPy array).\n        yolo_annotation: A string in YOLO segmentation format.\n        color: The color of the contour (B, G, R).\n        thickness: The thickness of the contour lines.\n\n    Returns:\n        The image with drawn contours.\n    \"\"\"\n    h, w, _ = image.shape\n    annotated_image = image.copy()\n    if not yolo_annotation:\n        return annotated_image\n\n    for points in yolo_annotation:\n        # Reshape to (n_points, 1, 2) for cv2.polylines\n        contour = np.array(points).reshape(-1, 2)\n        # Denormalize the coordinates\n        denormalized_contour = np.array([[int(p[0] * w), int(p[1] * h)] for p in contour])\n        # Draw the contour\n        cv2.polylines(annotated_image, [denormalized_contour], isClosed=True, color=color, thickness=thickness)\n\n    return annotated_image\n\n# Draw the YOLO segmentation on the sample image\nannotated_sample_image = draw_yolo_segmentation(sample_image, sample_yolo_annotation)\n\n# Display the annotated image and the YOLO annotation side by side\nplt.imshow(cv2.cvtColor(annotated_sample_image, cv2.COLOR_BGR2RGB))\nplt.title(\"Annotated Sample Image\")\nplt.axis('off')\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-29T08:26:57.028569Z","iopub.execute_input":"2025-10-29T08:26:57.028887Z","iopub.status.idle":"2025-10-29T08:26:57.223920Z","shell.execute_reply.started":"2025-10-29T08:26:57.028858Z","shell.execute_reply":"2025-10-29T08:26:57.223042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\nfrom tqdm import tqdm\n\n# Define output directories\nOUTPUT_DIR = '.'\n# OUTPUT_DIR = './yolo_dataset'\nos.makedirs(os.path.join(OUTPUT_DIR, 'images', 'train'), exist_ok=True)\nos.makedirs(os.path.join(OUTPUT_DIR, 'images', 'val'), exist_ok=True)\nos.makedirs(os.path.join(OUTPUT_DIR, 'labels', 'train'), exist_ok=True)\nos.makedirs(os.path.join(OUTPUT_DIR, 'labels', 'val'), exist_ok=True)\n\n# Paths to original data\ntrain_forged_images_dir = os.path.join(INPUT_DIR, 'train_images', 'forged')\ntrain_authentic_images_dir = os.path.join(INPUT_DIR, 'train_images', 'authentic') # Add path to authentic images\ntrain_masks_dir = os.path.join(INPUT_DIR, 'train_masks')\n\n# Get list of forged and authentic image names\nforged_image_names = [f for f in os.listdir(train_forged_images_dir) if f.endswith('.png')]\nauthentic_image_names = [f for f in os.listdir(train_authentic_images_dir) if f.endswith('.png')] # Get authentic image names\n\n# Split forged data into training and validation sets\ntrain_forged_image_names, val_forged_image_names = train_test_split(forged_image_names, test_size=0.2, random_state=42)\n\n# Split authentic data into training and validation sets\ntrain_authentic_image_names, val_authentic_image_names = train_test_split(authentic_image_names, test_size=0.2, random_state=42)\n\n\n# Prepare data for processing in a single loop\ndataset_info = []\n# Add forged images\nfor img_name in train_forged_image_names:\n    dataset_info.append((img_name, 'train', 'forged'))\nfor img_name in val_forged_image_names:\n    dataset_info.append((img_name, 'val', 'forged'))\n\n# Add authentic images\nfor img_name in train_authentic_image_names:\n    dataset_info.append((img_name, 'train', 'authentic'))\nfor img_name in val_authentic_image_names:\n    dataset_info.append((img_name, 'val', 'authentic'))\n\n\nprint(\"Processing dataset...\")\nfor image_name, dataset_type, image_type in tqdm(dataset_info):\n    if image_type == 'forged':\n        image_path = os.path.join(train_forged_images_dir, image_name)\n        mask_name = image_name.replace('.png', '.npy')\n        mask_path = os.path.join(train_masks_dir, mask_name)\n\n        if not os.path.exists(mask_path):\n            print(f\"Missing mask for {mask_path}\")\n            continue\n\n        mask = load_mask(mask_path)\n        yolo_annotation = mask_to_yolo_segmentation(mask, tolerance=0.005)\n\n        # Save annotation\n        label_filename = f'{image_type}_{image_name.replace(\".png\", \".txt\")}' # Add image_type prefix\n        label_path = os.path.join(OUTPUT_DIR, 'labels', dataset_type, label_filename)\n        with open(label_path, 'w') as fl:\n            fl.write(format_yolo_segmentation(yolo_annotation))\n\n    elif image_type == 'authentic':\n        image_path = os.path.join(train_authentic_images_dir, image_name)\n        # No mask or label file for authentic images\n\n    # Copy image to output directory\n    image = cv2.imread(image_path)\n    output_image_name = f'{image_type}_{image_name}' # Add image_type prefix\n    output_image_path = os.path.join(OUTPUT_DIR, 'images', dataset_type, output_image_name)\n    cv2.imwrite(output_image_path, image)\n\n\nprint(\"Dataset conversion to YOLO format complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-29T08:26:57.225741Z","iopub.execute_input":"2025-10-29T08:26:57.225998Z","iopub.status.idle":"2025-10-29T08:40:20.446004Z","shell.execute_reply.started":"2025-10-29T08:26:57.225977Z","shell.execute_reply":"2025-10-29T08:40:20.444544Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import yaml\n\n# Define the content of the dataset.yaml file\ndataset_yaml_content = {\n    # Absolute path when the generated dataset is added to your notebook\n    'path': \"/kaggle/input/forgerydetection-yolo-segmentation\",\n    'train': 'images/train',  # Relative path to the training images folder\n    'val': 'images/val',      # Relative path to the validation images folder\n    'nc': 1,                  # Number of classes\n    'names': ['forged']       # Class names\n}\n\n# Define the path to save the dataset.yaml file\ndataset_yaml_path = os.path.join(OUTPUT_DIR, 'dataset.yaml')\n\n# Write the content to the dataset.yaml file\nwith open(dataset_yaml_path, 'w') as f:\n    yaml.dump(dataset_yaml_content, f, default_flow_style=None)\n\nprint(f\"Dataset YAML file created at: {dataset_yaml_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-29T08:42:37.917349Z","iopub.execute_input":"2025-10-29T08:42:37.917715Z","iopub.status.idle":"2025-10-29T08:42:37.925673Z","shell.execute_reply.started":"2025-10-29T08:42:37.917688Z","shell.execute_reply":"2025-10-29T08:42:37.924762Z"}},"outputs":[],"execution_count":null}]}