{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":226382143,"sourceType":"kernelVersion"},{"sourceId":241460991,"sourceType":"kernelVersion"}],"dockerImageVersionId":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# BYU Locating Flagellar Motors\n\n## Data Preprocessing Pipeline\n\nThis is the first notebook in a series representing my submission for the BYU Locating Bacterial Flagellar Motors 2025 Kaggle challenge. The notebooks in this series cover the full workflow from data preprocessing to model training and submission generation.\n\n### Notebook Series:\n1. **Parse Data (Current)**: Extracting and preparing 2D slices containing motors to make a YOLO dataset\n         This is a version modified to do data augmentation (for vertical and horizontal flips)\n2. **Simple Data Augmentation (Current)**  Augment the yolo data by flipping the images (and labels) vertically and horizontally)\n3. **[Visualize Data](https://www.kaggle.com/code/andrewjdarley/visualize-data)**: Exploratory data analysis and visualization of annotated motor locations\n4. **[Train YOLO](https://www.kaggle.com/code/andrewjdarley/train-yolo)**: Fine tuning an YOLOv8 object detection model on the prepared dataset\n5. **[Submission Notebook](https://www.kaggle.com/code/andrewjdarley/submission-notebook)**: Running inference and generating submission files\n \n## About this Notebook\n\nThis notebook preprocesses the tomographic data for training. It performs the following:\n\n1. **Data Loading**: Reads the tomograms and motor location annotations\n2. **Slice Extraction**: Extracts 2D slices containing motors and surrounding slices (± TRUST parameter, little experimentation has been done modifying this parameter)\n3. **Image Normalization**: Normalizes slice intensity using percentile-based contrast enhancement (standard across all my work)\n4. **Dataset Preparation**: Organizes data into YOLO-compatible format with:\n   - Train/validation split of motors with no overlap in source tomograms (80/20 split)\n   - Bounding box annotations for each motor (I chose an arbitrary box size. It could be optimized)\n   - Proper directory structure for YOLO training\n5. **Configuration**: Generates dataset.yaml\n\nMuch of the code was improved by github-copilot.","metadata":{}},{"cell_type":"code","source":"import glob\nimport os\n\nimport numpy as np\nimport pandas as pd\nfrom typing import List, Optional\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport time\n\nfrom PIL import Image, ImageDraw","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:13.798720Z","iopub.execute_input":"2025-05-23T17:30:13.799168Z","iopub.status.idle":"2025-05-23T17:30:13.805190Z","shell.execute_reply.started":"2025-05-23T17:30:13.799128Z","shell.execute_reply":"2025-05-23T17:30:13.804058Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#%env JOBLIB_TEMP_FOLDER=/tmp\n# this may be necessary; because I was getting 'out of space' errors.","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:13.807216Z","iopub.execute_input":"2025-05-23T17:30:13.807525Z","iopub.status.idle":"2025-05-23T17:30:13.826385Z","shell.execute_reply.started":"2025-05-23T17:30:13.807495Z","shell.execute_reply":"2025-05-23T17:30:13.825243Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"my_data_version = 14","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:13.827856Z","iopub.execute_input":"2025-05-23T17:30:13.828194Z","iopub.status.idle":"2025-05-23T17:30:13.848503Z","shell.execute_reply.started":"2025-05-23T17:30:13.828152Z","shell.execute_reply":"2025-05-23T17:30:13.847264Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Classes for augmenting data\n\nKeeps labels and yolo image data in sync.\n\nTraining the yolo model uses .jpg and .txt files;  the 'truth' for the .txt files comes from labels.csv\nTo train with modified .jpgs, you need to modify the .txt files\nTo score with modified .jpgs, you need to modify a copy of the labels.csv.  (it would also desirable to transform the output answers, instead; but I didn't do that.","metadata":{}},{"cell_type":"code","source":"class SliceAndTxt:\n    def __init__(self, slice_base_dir: str, label_dir: Optional[str], slice_name: str):\n        \"\"\"\n        Initialize the SliceAndTxt class.\n\n        :param slice_base_dir: Directory containing the slice images.\n        :param label_dir: Directory containing the labels (optional).\n        :param slice_name: Name of the slice.\n        \"\"\"\n        self.slice_base_dir = slice_base_dir\n        self.label_dir = label_dir\n        self.slice_name = slice_name\n        self.img: Optional[Image.Image] = None\n        # self.exif: Optional[dict] = None\n        self.annotations_df: Optional[pd.DataFrame] = None\n        self.output_tag = ''\n\n    def read(self, do_print = False) -> None:\n        \"\"\"\n        Read the slice image and its annotations (if label_dir is provided).\n        \"\"\"\n        slice_file = os.path.join(self.slice_base_dir, self.slice_name) + '.jpg'\n\n        if do_print:\n            print (f'{slice_file=}')\n\n        self.tomo_name, self.z,self.y,self.x = self.get_coords_from_filename()\n        \n        try:\n            self.img = Image.open(slice_file)\n            # self.exif = self.img.getexif()\n        except FileNotFoundError:\n            print(f\"Error: File not found - {slice_file}\")\n        except Exception as e:\n            print(f\"Error reading image: {e}\")\n            \n        self.read_annotations(do_print = do_print)\n    \n    def read_annotations(self, do_print = False) -> None:\n        \"\"\"\n        Read the annotations for the slice (if label_dir is provided).\n        \"\"\"\n        if not self.label_dir:\n            return\n\n        label_file = os.path.join(self.label_dir, self.slice_name) + '.txt'\n        if do_print:\n            print (f'{label_file=}')\n        annotations = []\n        img_width, img_height = self.img.size\n\n        if os.path.exists(label_file):\n            with open(label_file, 'r') as f:\n                \n                tomo, z,y,x = self.get_coords_from_filename()\n                for line in f:\n                    values = line.strip().split() # 0 x_center_norm y_center_norm, x_width_norm, y_width_norm\n                    class_id = int(values[0])\n\n                    _, x_center_norm, y_center_norm, width, height = values\n                    x_center_norm = float(x_center_norm)\n                    y_center_norm = float(y_center_norm)\n                    x2 = x_center_norm * img_width\n                    y2 = y_center_norm * img_height\n\n                    assert abs(x-x2) < 0.0001, f'{x=}, {x2=}'\n                    assert abs(y-y2) < 0.0001, f'{y=}, {y2=}'\n\n                    annotations.append({\n                        'class_id': class_id,\n                        'width': width,\n                        'height': height,\n                        'x':x, \n                        'y':y, \n                        'z':z\n                    })\n            self.annotations_df = pd.DataFrame(annotations)\n        else:\n            self.annotations_df = pd.DataFrame()\n\n    def get_coords_from_filename(self):\n\n        name_coords = self.slice_name.split(\"tomo_\")\n        coords = name_coords[1].split('_')\n        tomo = coords[0]\n        z = int(coords[1][1:])\n        y = int(coords[2][1:])\n        x = int(coords[3][1:5])\n\n        return tomo, z, y, x \n\n    def output_slice_name(self):\n        output_name = f'tomo_{self.tomo_name}_z{self.z:04d}_y{self.y:04d}_x{self.x:04d}{self.output_tag}'\n        return output_name\n    \n\n    def write(self, out_slice_base_dir, out_label_base_dir,  do_print: bool = False) -> None:\n        \"\"\"\n        Write the slice image and its annotations (if label_dir is provided).\n\n        :param new_name: New name for the slice image and annotations.\n        :param do_print: Whether to print debug information.\n        \"\"\"\n        output_name = f'{out_slice_base_dir}/{self.output_slice_name()}.jpg'\n\n        os.makedirs(os.path.dirname(output_name), exist_ok = True)\n        self.img.save(output_name)\n        if do_print:\n            print(f\"Wrote image of size {self.img.size} to {output_name}\")\n        if self.label_dir:\n            output_label_name = output_name.replace('image', 'label').replace('.jpg', '.txt')\n\n            if do_print:\n                print(f'{output_name=}\\n{output_label_name=}')\n            self.write_annotations(output_label_name,  do_print)\n\n    def write_annotations(self, output_label_name,  do_print: bool = False) -> None:\n        \"\"\"\n        Write the annotations for the slice (if label_dir is provided).\n\n        :param do_print: Whether to print debug information.\n        \"\"\"\n        if not self.label_dir:\n            return\n\n        img_width, img_height = self.img.size\n\n        df = self.annotations_df.copy()\n        df['class'] = 0\n        df['x_center_norm'] = df['x'] / img_width\n        df['y_center_norm'] = df['y'] / img_height\n\n        out_df = df[['class', 'x_center_norm', 'y_center_norm', 'width', 'height']]\n         \n        os.makedirs(os.path.dirname(output_label_name), exist_ok = True)\n        out_df.to_csv(output_label_name, sep=' ', header=False, index=False)\n        if do_print:\n            print(f\"Wrote {out_df.shape[0]} lines to {output_label_name}\")\n\n    def print(self, print_details: bool = True) -> None:\n        \"\"\"\n        Print the details of the slice and its annotations.\n\n        :param print_details: Whether to print detailed information.\n        \"\"\"\n        print(f\"{self.slice_base_dir=}\\n{self.label_dir=}\\n{self.slice_name=}\\n{self.output_tag=}\")\n    \n        if print_details and self.annotations_df is not None:\n            print(f\"\\n{self.annotations_df=}\")\n    \n    def show(self) -> None:\n        \"\"\"\n        Show the slice image with annotations (if label_dir is provided).\n        \"\"\"\n        if  self.annotations_df is None:\n            self.img.show()\n            return\n\n        num_samples = 1\n        rows = int(np.ceil(num_samples / 2))\n        cols = min(num_samples, 2)\n        fig, axes = plt.subplots(rows, cols, figsize=(14, 5 * rows))\n        \n        # Handle the case of a single subplot\n        if num_samples == 1:\n            axes = np.array([axes])\n    \n        # Flatten axes array for easy indexing\n        axes = axes.flatten()\n\n        overlay = Image.new('RGBA', self.img.size, (0, 0, 0, 0))\n        draw = ImageDraw.Draw(overlay)\n        img_width, img_height = self.img.size\n\n        for _, ann in self.annotations_df.iterrows():\n            x = float(ann['x'])\n            y = float(ann['y'])\n            width = 24\n            height = 24\n            x1 = max(0, int(x - width / 2))\n            y1 = max(0, int(y - height / 2))\n            x2 = min(img_width, int(x + width / 2))\n            y2 = min(img_height, int(y + height / 2))\n            draw.rectangle([x1, y1, x2, y2], outline=(255, 0, 0, 200))\n\n        annotated_img = Image.alpha_composite(self.img.convert('RGBA'), overlay).convert('RGB')\n        axes[0].imshow(np.array(annotated_img))\n\n        plt.show()\n\n    def flip_horizontal(self) -> None:\n        \"\"\"\n        Flip the slice image and annotations horizontally (if label_dir is provided).\n        \"\"\"\n        img_width, _ = self.img.size\n        self.img = self.img.transpose(Image.FLIP_LEFT_RIGHT)\n        if self.label_dir and self.annotations_df is not None:\n\n            # self.annotations_df['x_center'] = img_width - self.annotations_df['x_center']\n            # self.annotations_df['x_center_norm'] = 1 - self.annotations_df['x_center_norm']\n            self.annotations_df['x'] = img_width - self.annotations_df['x']\n            self.x = img_width - self.x\n            \n        self.output_tag += '-fh'\n\n    def flip_vertical(self) -> None:\n        \"\"\"\n        Flip the slice image and annotations vertically (if label_dir is provided).\n        \"\"\"\n        _, img_height = self.img.size\n        self.img = self.img.transpose(Image.FLIP_TOP_BOTTOM)\n        if self.label_dir and self.annotations_df is not None:            \n            # self.annotations_df['y_center'] = img_height - self.annotations_df['y_center']\n            # self.annotations_df['y_center_norm'] = 1 - self.annotations_df['y_center_norm']\n            self.annotations_df['y'] = img_height - self.annotations_df['y']\n            self.y = img_height - self.y  # used for the filename\n            \n        self.output_tag += '-fv'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:13.850298Z","iopub.execute_input":"2025-05-23T17:30:13.850652Z","iopub.status.idle":"2025-05-23T17:30:13.878863Z","shell.execute_reply.started":"2025-05-23T17:30:13.850624Z","shell.execute_reply":"2025-05-23T17:30:13.877794Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Sanity check on SliceAndTxt class\n\nCheck visually that the image is transformed, and that the red bounding box is adjusted appropriately.\n\n","metadata":{}},{"cell_type":"code","source":"my_data_version = 14","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:13.880255Z","iopub.execute_input":"2025-05-23T17:30:13.880615Z","iopub.status.idle":"2025-05-23T17:30:13.899955Z","shell.execute_reply.started":"2025-05-23T17:30:13.880573Z","shell.execute_reply":"2025-05-23T17:30:13.898740Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!mkdir -p tmp/image tmp/label\nbase_dir = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/images/train'\nlabel_dir = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/labels/train'\nout_dir = 'tmp/image'\nout_label_dir = 'tmp/label'\n\nslice_name = os.listdir(base_dir)[40].replace('.jpg', '')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:13.901245Z","iopub.execute_input":"2025-05-23T17:30:13.901682Z","iopub.status.idle":"2025-05-23T17:30:14.165731Z","shell.execute_reply.started":"2025-05-23T17:30:13.901625Z","shell.execute_reply":"2025-05-23T17:30:14.164140Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ls /kaggle/input/parse-data/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:14.167219Z","iopub.execute_input":"2025-05-23T17:30:14.167623Z","iopub.status.idle":"2025-05-23T17:30:14.294399Z","shell.execute_reply.started":"2025-05-23T17:30:14.167575Z","shell.execute_reply":"2025-05-23T17:30:14.293166Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#### Check that reading and writing, without transforming, works ok","metadata":{}},{"cell_type":"markdown","source":"###","metadata":{}},{"cell_type":"code","source":"s1 = SliceAndTxt(base_dir, label_dir, slice_name)\ns1.read()\ns1.write(out_dir, out_label_dir)\ns1.show()\ns1.print()\n\nprint ('on to s2')\ns2 = SliceAndTxt(out_dir, out_label_dir, slice_name)\ns2.read()\ns2.show()\ns2.print()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:14.295627Z","iopub.execute_input":"2025-05-23T17:30:14.295917Z","iopub.status.idle":"2025-05-23T17:30:15.344576Z","shell.execute_reply.started":"2025-05-23T17:30:14.295890Z","shell.execute_reply":"2025-05-23T17:30:15.343260Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#### Check that horizontal flip works OK.","metadata":{}},{"cell_type":"code","source":"# check that horizontal flip works ok\ns3 = SliceAndTxt(base_dir, label_dir, slice_name)\ns3.read()\ns3.flip_horizontal()\ns3.write(out_dir, out_label_dir)\nfh_slice_name = s3.output_slice_name()\ns3.show()\ns3.print()\n\n\nprint ('on to s3b')\n\ns3b = SliceAndTxt(out_dir, out_label_dir,  fh_slice_name)\ns3b.read()\n\ns3b.show()\ns3b.print()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:15.345853Z","iopub.execute_input":"2025-05-23T17:30:15.346270Z","iopub.status.idle":"2025-05-23T17:30:16.082568Z","shell.execute_reply.started":"2025-05-23T17:30:15.346232Z","shell.execute_reply":"2025-05-23T17:30:16.081251Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#### Check vertical flip","metadata":{}},{"cell_type":"code","source":"# check vertical flip\n\ns4 = SliceAndTxt(base_dir, label_dir, slice_name)\ns4.read()\ns4.show()\ns4,print()\ns4.flip_vertical()\ns4.write(out_dir, out_label_dir)\nfh_slice_name = s4.output_slice_name()\ns4.show()\ns4.print()\n\n\ns4b = SliceAndTxt(out_dir, out_label_dir,  fh_slice_name)\ns4b.read()\n\ns4b.show()\ns4b.print()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:16.086568Z","iopub.execute_input":"2025-05-23T17:30:16.087065Z","iopub.status.idle":"2025-05-23T17:30:17.298505Z","shell.execute_reply.started":"2025-05-23T17:30:16.086987Z","shell.execute_reply":"2025-05-23T17:30:17.297106Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Class to applySliceAndTxt to whole tomograms.","metadata":{}},{"cell_type":"code","source":"# class for handling all the slices in a tomogram\n\nclass TomogramAndTxts:\n    def __init__(self, base_dir: str, label_dir: Optional[str], tomogram_name: str, do_print = False):\n        \"\"\"\n        Initialize the TomogramAndTxts class.\n\n        :param base_dir: Directory containing the tomogram images.\n        :param label_dir: Directory containing the labels (optional).\n        :param tomogram_name: Name of the tomogram.\n        \"\"\"\n        self.base_dir = base_dir\n        self.label_dir = label_dir\n        self.tomogram_name = tomogram_name\n        self.slices: List[SliceAndTxt] = []\n\n        # Find slice names\n        glob_str = f'{self.base_dir}/tomo_{self.tomogram_name}_*.jpg'\n        slice_full_names = glob.glob(glob_str)\n\n        if do_print:\n            print (f'{glob_str=}\\n{slice_full_names=}')\n        self.slice_names = [os.path.basename(slice_full_name).replace('.jpg', '') for slice_full_name in slice_full_names]\n\n        # Initialize slices\n        self.slices = [SliceAndTxt(self.base_dir, self.label_dir, slice_name,) for slice_name in self.slice_names]\n\n        # Read slices\n        for slice in self.slices:\n            slice.read()\n\n    def write(self, out_slice_base_dir, out_label_base_dir, do_print: bool = False) -> None:\n        \"\"\"\n        Write the tomogram slices and their labels (if label_dir is provided).\n\n        :param do_print: Whether to print debug information.\n        \"\"\"\n        for slice in self.slices:\n            slice.write(out_slice_base_dir, out_label_base_dir, do_print=do_print)\n        if do_print:\n            print(f\"Wrote tomogram with {len(self.slices)} slices and {self.count_motors()} motors.\")\n\n    def flip_horizontal(self) -> None:\n        \"\"\"\n        Flip the tomogram slices horizontally.\n        \"\"\"\n        self._apply_to_slices(\"flip_horizontal\")\n\n    def flip_vertical(self) -> None:\n        \"\"\"\n        Flip the tomogram slices vertically.\n        \"\"\"\n        self._apply_to_slices(\"flip_vertical\")\n\n    def show(self) -> None:\n        \"\"\"\n        Show the tomogram slices.\n        \"\"\"\n        self._apply_to_slices(\"show\")\n\n    def print(self) -> None:\n        \"\"\"\n        Print the details of the tomogram and its slices.\n        \"\"\"\n        print(f\"{self.base_dir=}\\n{self.label_dir=}\\n{self.tomogram_name=}\\n\")\n        print(f\"Total motors = {self.count_motors()}\")\n\n    def count_motors(self) -> int:\n        \"\"\"\n        Count the total number of motors in the tomogram.  /this is quite wrong...\n\n        :return: Total number of motors.\n        \"\"\"\n        total_motors = 0\n        for slice in self.slices:\n            if slice.annotations_df is not None:\n                total_motors += slice.annotations_df.shape[0]\n        return total_motors\n\n    def _apply_to_slices(self, method_name: str) -> None:\n        \"\"\"\n        Apply a method to all slices in the tomogram.\n\n        :param method_name: Name of the method to apply.\n        \"\"\"\n        for slice in self.slices:\n            method = getattr(slice, method_name, None)\n            if callable(method):\n                method()\n            else:\n                print(f\"Method {method_name} not found in SliceAndTxt.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:17.300021Z","iopub.execute_input":"2025-05-23T17:30:17.300341Z","iopub.status.idle":"2025-05-23T17:30:17.312133Z","shell.execute_reply.started":"2025-05-23T17:30:17.300311Z","shell.execute_reply":"2025-05-23T17:30:17.310950Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndata_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\ndef augmentation_vertical_and_horizontal (data_path = data_path, \n                                          image_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/images/train', \n                                          label_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/labels/train',\n                                          output_image_path = f'/kaggle/working/yolo_dataset_{my_data_version}/images/train',\n                                          output_label_path = f'/kaggle/working/yolo_dataset_{my_data_version}/labels/train',\n                                          do_print = False,\n                                          N_tomos = -1,\n                                         ):\n\n    start_time = time.time()\n\n    labels_df = pd.read_csv(os.path.join(data_path, \"train_labels.csv\"))\n    tomo_df = labels_df[labels_df['Number of motors'] > 0].copy()\n    unique_tomos = tomo_df['tomo_id'].str.replace('tomo_', '').unique()\n    \n    for tomo in tqdm(unique_tomos if N_tomos == -1 else unique_tomos[:N_tomos]):\n\n        tomo_v = TomogramAndTxts(image_path, label_path,  tomo, do_print = do_print )\n        tomo_v.write(output_image_path, output_label_path, do_print = do_print)\n        tomo_v.flip_vertical()\n        tomo_v.write(output_image_path, output_label_path, do_print = do_print)\n\n        if do_print:\n            print (f'{tomo=}')\n        t = TomogramAndTxts(image_path, label_path,  tomo, do_print = do_print)\n        # print (f'{t.count_motors()=}')\n        t.write(output_image_path, output_label_path, do_print = do_print)\n        t.flip_horizontal()\n        t.write(output_image_path, output_label_path, do_print = do_print)\n        t.flip_vertical()\n        t.write(output_image_path, output_label_path, do_print = do_print)\n\n    elapsed = time.time() - start_time\n    print(f\"\\nTotal execution time: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:17.313286Z","iopub.execute_input":"2025-05-23T17:30:17.313633Z","iopub.status.idle":"2025-05-23T17:30:17.335538Z","shell.execute_reply.started":"2025-05-23T17:30:17.313602Z","shell.execute_reply":"2025-05-23T17:30:17.334143Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"labels = pd.read_csv(f'{data_path}/train_labels.csv')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:17.336661Z","iopub.execute_input":"2025-05-23T17:30:17.336999Z","iopub.status.idle":"2025-05-23T17:30:17.368340Z","shell.execute_reply.started":"2025-05-23T17:30:17.336969Z","shell.execute_reply":"2025-05-23T17:30:17.367054Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndef write_labels(labels_df, output_path, labels_df_columns, suffix ):\n    out_df = labels_df[labels_df_columns]\n    output_filename = os.path.join(output_path, f'train_labels{suffix}.csv') \n    out_df.to_csv(output_filename, index=False)\n    print(f'wrote {out_df.shape} to {output_filename}')\n    \ndef labels_flip_vertical( df):\n    df = df.copy()\n    df['Motor axis 1'] = df['height'] - df['Motor axis 1']  \n    df.loc[df['Number of motors'] == 0, 'Motor axis 1']  = -1\n    assert min(df['Motor axis 1'].values) >= -1 \n    return df\n\ndef labels_flip_horizontal( df):\n    df = df.copy()\n    df['Motor axis 2'] = df['width'] - df['Motor axis 2']\n    df.loc[df['Number of motors'] == 0, 'Motor axis 2'] = -1\n    assert min(df['Motor axis 2'].values) >= -1 \n    return df\n\ndef get_sizes(tomo, image_path = '/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/train/', do_print = False):\n\n    if do_print:\n        print(f'{image_path=}')\n\n    img_files = glob.glob(f'{image_path}/*{tomo}*/slice*.jpg')\n    if do_print:\n        print(f'{img_files=}')\n\n    if len(img_files) == 0:\n        print(f'no img_files found for {image_path}/*{tomo}*')\n        return None\n\n    img_file = img_files[0]\n\n    assert img_file.endswith('.jpg')\n\n    name_coords = img_file.split(\"tomo_\")\n    coords = name_coords[1].split('_')\n    tomo_id = f'tomo_{coords[0].split(\"/\")[0]}'\n    \n    if do_print:\n        print (f'{tomo=}, {tomo_id=},{img_file=}, {name_coords=}')\n    \n\n    img = Image.open(img_file)\n    img_width, img_height = img.size\n    out_df = pd.DataFrame({'tomo_id':tomo_id, 'height':img_height, 'width':img_width}, index=[0])#.drop_duplicates()\n\n    if do_print:\n        print (f'{tomo_id=}, {img_file=}, {name_coords=}, {img_width=}, {img_height=}, {out_df=}')\n    return out_df    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:17.369453Z","iopub.execute_input":"2025-05-23T17:30:17.369785Z","iopub.status.idle":"2025-05-23T17:30:17.383080Z","shell.execute_reply.started":"2025-05-23T17:30:17.369755Z","shell.execute_reply":"2025-05-23T17:30:17.381873Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pandas as pd\n\nos.makedirs('/kaggle/working/labels', exist_ok = True)\n\ndata_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\ndef augmentation_vertical_and_horizontal_label_df (data_path = data_path, \n                                                   image_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/images/train', \n                                                   output_path = '/kaggle/working/labels/',                                                   \n                                                   do_print = False,\n                                                   ):\n\n    labels_df = pd.read_csv(os.path.join(data_path, \"train_labels.csv\"))\n    labels_df_columns = labels_df.columns\n    \n    unique_tomos = labels_df['tomo_id'].drop_duplicates()\n\n    if do_print:\n        print(f'{unique_tomos=}')\n\n    size_list = []\n    \n    for tomo in tqdm(unique_tomos):\n        sizes = get_sizes(tomo)\n        size_list.append(sizes)\n\n    sizes_df = pd.concat(size_list)\n    \n    if do_print:\n        print(f'{len(size_list)=}\\n{sizes_df.describe()=}, {unique_tomos.shape=}')\n        print(f'{sizes_df=}, {sizes_df[\"tomo_id\"].head()=}, {labels_df[\"tomo_id\"].head()=}')\n\n    labels_df_w_sizes = labels_df.merge(sizes_df, on = 'tomo_id', validate = 'm:1')\n    assert labels_df_w_sizes.shape[0] == labels_df.shape[0], f\"{labels_df_w_sizes.shape=}, {labels_df.shape=}, {sizes_df.shape=}\"\n\n    write_labels(labels_df, output_path=output_path, labels_df_columns=labels_df_columns, suffix = '-plain')\n\n    labels_df_horizontal = labels_flip_horizontal(labels_df_w_sizes.copy())\n    write_labels(labels_df_horizontal, output_path=output_path, labels_df_columns=labels_df_columns, suffix = '-fh')\n\n    labels_df_vertical = labels_flip_vertical(labels_df_w_sizes.copy())\n    write_labels(labels_df_vertical, output_path=output_path, labels_df_columns=labels_df_columns, suffix = '-fv')\n\n    labels_df_horizontal_and_vertical = labels_flip_horizontal(labels_df_vertical)\n    write_labels(labels_df_horizontal_and_vertical, output_path=output_path, labels_df_columns=labels_df_columns, suffix = '-fh-fv')\n\naugmentation_vertical_and_horizontal_label_df()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:17.384062Z","iopub.execute_input":"2025-05-23T17:30:17.384473Z","iopub.status.idle":"2025-05-23T17:30:37.003189Z","shell.execute_reply.started":"2025-05-23T17:30:17.384442Z","shell.execute_reply":"2025-05-23T17:30:37.002215Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# augment the training data\n\naugmentation_vertical_and_horizontal (image_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/images/train', \n                                      label_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/labels/train', \n                                      output_image_path = f'/kaggle/working/yolo_dataset_{my_data_version}/images/train',\n                                      output_label_path = f'/kaggle/working/yolo_dataset_{my_data_version}/labels/train',\n                                     )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:30:37.004421Z","iopub.execute_input":"2025-05-23T17:30:37.004746Z","iopub.status.idle":"2025-05-23T17:31:38.108050Z","shell.execute_reply.started":"2025-05-23T17:30:37.004710Z","shell.execute_reply":"2025-05-23T17:31:38.106312Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#augment the validationg data\naugmentation_vertical_and_horizontal (image_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/images/val', \n                                      label_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/labels/val', \n                                      output_image_path = f'/kaggle/working/yolo_dataset_{my_data_version}/images/val',\n                                      output_label_path = f'/kaggle/working/yolo_dataset_{my_data_version}/labels/val',\n                                     )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:31:38.108899Z","iopub.status.idle":"2025-05-23T17:31:38.109345Z","shell.execute_reply":"2025-05-23T17:31:38.109180Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\ndef show_augmentation_vertical_and_horizontal (data_path = data_path, \n                                          # image_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/images/train', \n                                          # label_path = f'/kaggle/input/parse-data/yolo_dataset_{my_data_version}/labels/train',\n                                          output_image_path = f'/kaggle/working/yolo_dataset_{my_data_version}/images/train',\n                                          output_label_path = f'/kaggle/working/yolo_dataset_{my_data_version}/labels/train',\n                                          do_print = True,\n                                          N_tomos = 3):\n\n    start_time = time.time()\n\n    labels_df = pd.read_csv(os.path.join(data_path, \"train_labels.csv\"))\n    tomo_df = labels_df[labels_df['Number of motors'] > 0].copy()\n    unique_tomos = tomo_df['tomo_id'].str.replace('tomo_', '').unique()\n\n    print(f'{len(unique_tomos)=}')\n    \n    for tomo in tqdm(unique_tomos if N_tomos == -1 else unique_tomos[:N_tomos]):\n        for variant in ['_fh', '_plain']:\n            if do_print:\n                print (f'{tomo=}, {variant=}')\n                \n            t = TomogramAndTxts(image_path+variant, label_path+variant,  tomo, do_print = do_print)\n            # print (f'{t.count_motors()=}')\n            t.show()\n\n    elapsed = time.time() - start_time\n    print(f\"\\nTotal execution time: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:31:38.110083Z","iopub.status.idle":"2025-05-23T17:31:38.110487Z","shell.execute_reply":"2025-05-23T17:31:38.110316Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Check that the 'plain' transformation  makes files matching the original","metadata":{}},{"cell_type":"code","source":"import filecmp\nimport os\n\ndef compare_directories_recursive(dir1, dir2, ignore_file_set = set()):\n    non_equals = 0\n    for root1, dirs1, files1 in os.walk(dir1):\n        root2 = root1.replace(dir1, dir2, 1)  # Map corresponding directory in dir2\n\n        if not os.path.exists(root2):\n            print(f\"Directory missing: {root2}\")\n            return False\n\n        # Compare directories\n        dirs2 = set(os.listdir(root2)) - set(files1)  # Filter out files\n        dirs1_set = set(dirs1)\n\n        if dirs1_set != dirs2:\n            print(f\"Different subdirectories in: {root1} vs {root2}\\n{dirs1_set=}\\n{dirs2=}\")\n            return False\n\n        # Compare files\n        files2 = set(os.listdir(root2)) - set(dirs1)  # Filter out directories\n        files1_set = set(files1)\n\n        if (files1_set - ignore_file_set) != (files2 - ignore_file_set):\n            print(f\"Different files in: {root1} vs {root2}\")\n            print(f'{files1_set - files2=}')\n            print(f'{files2     - files1_set=}')\n            intersect = files1_set.intersection(files2)\n            print(f'{list(intersect)[:5]=}\\n{list(intersect)[-5:]=}')\n            return False\n\n        for file in files1:\n            path1 = os.path.join(root1, file)\n            path2 = os.path.join(root2, file)\n\n            if not filecmp.cmp(path1, path2, shallow=False):\n                print(f\"File content differs: {path1} vs {path2}\")\n                non_equals += 1\n\n    if non_equals == 0:\n        print(\"Directories and all subdirectories are identical.\")\n        return True\n    else:\n        print(f'{non_equals} files are different')\n        return False\n\n# Example usage\ndir1 = \"path/to/directory1\"\ndir2 = \"path/to/directory2\"\ncompare_directories_recursive(dir1, dir2, set())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:31:38.111452Z","iopub.status.idle":"2025-05-23T17:31:38.111931Z","shell.execute_reply":"2025-05-23T17:31:38.111706Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/working/yolo_dataset_{my_data_version} /tmp\n\n!rm /tmp/yolo_dataset_{my_data_version}/*/*/*-fh*\n!rm /tmp/yolo_dataset_{my_data_version}/*/*/*-fv*","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:31:38.112613Z","iopub.status.idle":"2025-05-23T17:31:38.112894Z","shell.execute_reply":"2025-05-23T17:31:38.112777Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"rm -rf /kaggle/working/yolo_dataset_{my_data_version}/yolo_dataset.andrew_darling/yolo_dataset","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:31:38.113517Z","iopub.status.idle":"2025-05-23T17:31:38.113898Z","shell.execute_reply":"2025-05-23T17:31:38.113717Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dir1 = \"/kaggle/input/parse-data/yolo_dataset/labels\"\ndir2 = f\"/tmp/yolo_dataset_{my_data_version}/labels\"\nassert compare_directories_recursive(dir1, dir2)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-23T17:31:38.114674Z","iopub.status.idle":"2025-05-23T17:31:38.114991Z","shell.execute_reply":"2025-05-23T17:31:38.114858Z"}},"outputs":[],"execution_count":null}]}