{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":91249,"databundleVersionId":11294684},{"sourceType":"datasetVersion","sourceId":11697608,"datasetId":7206328,"databundleVersionId":12173114},{"sourceType":"datasetVersion","sourceId":11646383,"datasetId":7308315,"databundleVersionId":12114754},{"sourceType":"datasetVersion","sourceId":11495557,"datasetId":7206260,"databundleVersionId":11942833},{"sourceType":"datasetVersion","sourceId":11640450,"datasetId":7304147,"databundleVersionId":12107893}],"dockerImageVersionId":31011,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip install --no-index --find-links=/kaggle/input/ultralytics-offline-install2/ultralytics_pk \\\n#     torch torchvision torchaudio ultralytics \n# import kagglehub\n# kagglehub.dataset_download('rachiteagles/yolo-pkg')\n# kagglehub.dataset_download('rachiteagles/yolo-model')\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-07T12:09:58.057496Z","iopub.execute_input":"2025-05-07T12:09:58.057676Z","iopub.status.idle":"2025-05-07T12:09:58.062025Z","shell.execute_reply.started":"2025-05-07T12:09:58.057660Z","shell.execute_reply":"2025-05-07T12:09:58.061352Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --no-index --no-deps /kaggle/input/yolo-pkg/yolo/ultralytics-8.3.112-py3-none-any.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T12:09:58.063532Z","iopub.execute_input":"2025-05-07T12:09:58.063754Z","iopub.status.idle":"2025-05-07T12:10:01.003978Z","shell.execute_reply.started":"2025-05-07T12:09:58.063733Z","shell.execute_reply":"2025-05-07T12:10:01.003024Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom matplotlib.patches import Rectangle\nimport os\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\nfrom ultralytics import YOLO\nimport cv2\nimport yaml\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nimport os\nos.environ[\"ULTRALYTICS_CALLBACKS_DISABLE\"] = \"raytune\"\n\nroot_path = '/kaggle/input/byu-locating-bacterial-flagellar-motors-2025'\ntrain_path = os.path.join(root_path,'train')\ntest_path = os.path.join(root_path,'test')\ntrain_labels_path= os.path.join(root_path,'train_labels.csv')\ntrain_labels = pd.read_csv(train_labels_path)\n\nTrust = 4\nsplit_ratio=0.8\nBOX_SIZE=24","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T12:10:01.005123Z","iopub.execute_input":"2025-05-07T12:10:01.005413Z","iopub.status.idle":"2025-05-07T12:10:08.357783Z","shell.execute_reply.started":"2025-05-07T12:10:01.005382Z","shell.execute_reply":"2025-05-07T12:10:08.357063Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Preparing YOLO dataset","metadata":{}},{"cell_type":"code","source":"'''\n    1. create a folder for the yolo dataset both for training and validation\n    2. extract the tomo_id that have motors\n    3. split the ids into training and validation where split ratio is 0.8\n    4. for each motor collect (2*Trust + 1)slices , normalize and save to the yolo_dataset\n    5. create the yaml configuration\n'''\n\n# Step1\n\nyolo_data_path = '/kaggle/working/'\nyolo_train_img = os.path.join(yolo_data_path,'images','train')\nyolo_train_label= os.path.join(yolo_data_path,'labels','train')\nyolo_val_img = os.path.join(yolo_data_path,'images','val')\nyolo_val_label=os.path.join(yolo_data_path ,'labels','val')\nfor path in [yolo_train_img,yolo_train_label,yolo_val_img,yolo_val_label]:\n    os.makedirs(path, exist_ok=True)\n\n#steps 2 and 3\n\nmotor_tomo_id = train_labels[train_labels['Number of motors'] > 0]['tomo_id'].unique()\nnp.random.shuffle(motor_tomo_id)\ntrain_ids = motor_tomo_id[:int(len(motor_tomo_id)*split_ratio)]\nval_ids = motor_tomo_id[int(len(motor_tomo_id)*split_ratio):]\n\n# Steps 4\n\ndef normalize_slice(img):\n    p2=np.percentile(img,2)\n    p98=np.percentile(img,98)\n    clipped_data = np.clip(img,p2,p98)\n    normalize_img =  255*(img-p2)/(p98-92)\n    return np.uint8(normalize_img)\n\ndef process_yolo_dataset(train_ids,yolo_train_img,yolo_train_label):\n    col_name=['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2','Array shape (axis 0)']\n    new_col_name = ['tomo_id','z_axis','y_axis','x_axis','z_max']\n    train_df = train_labels[train_labels['tomo_id'].isin(train_ids)][col_name]\n    train_df = train_df.sort_values(by='tomo_id').reset_index(drop=True)\n    print('total motors',len(train_df))\n    train_df.columns = new_col_name\n\n    process_slice=0\n    for motor in tqdm(train_df.itertuples(index=False),desc=f'processing slices'):\n        tomo_id,z_centre,y_centre,x_centre,max_slice = motor.tomo_id , motor.z_axis,motor.y_axis,motor.x_axis,motor.z_max\n        z_min ,z_max= int(max(0,z_centre-Trust)),int(min(z_centre + Trust , max_slice))\n        for z in range(int(z_min),int(z_max+1)):\n            slice_name =f'slice_{z:04d}.jpg'\n            slice_path = os.path .join(train_path,tomo_id,slice_name)\n            img= Image.open(slice_path)\n            img_arr= np.array(img)\n            img_normalize = normalize_slice(img_arr)\n            file_name= f\"{tomo_id}_z{int(z):04d}_y{int(y_centre):04d}_x{int(x_centre):04d}.jpg\"\n            file_path = os.path.join(yolo_train_img,file_name)\n            Image.fromarray(img_arr).save(file_path)\n            img_width,img_height=img.size\n            x_centre_norm= x_centre/img_width\n            y_centre_norm=y_centre/img_height\n            box_width = BOX_SIZE/img_width\n            box_height=BOX_SIZE/img_height\n            label_path = os.path.join(yolo_train_label,file_name.replace('.jpg','.txt'))\n            process_slice+=1\n            with open(label_path,'w') as f:\n                f.write(f\"0 {x_centre_norm} {y_centre_norm} {box_width} {box_height}\\n\")\n\n    return process_slice\n\ntrain_slice = process_yolo_dataset(train_ids,yolo_train_img,yolo_train_label)\nprint(f'total slice for training  is {train_slice}')\nval_slice = process_yolo_dataset(val_ids,yolo_val_img,yolo_val_label)\nprint(f'total slice for validation  is {val_slice}')\n\n# step5\n\nyaml_content={\n        'path':yolo_data_path,\n        'train':'images/train',\n        'val':'images/train' if not os.path.exists(yolo_val_img) else 'images/val',\n        'names':{0:'motor'}\n    }\nwith open(os.path.join(yolo_data_path,'dataset.yaml') ,'w') as f:\n    yaml.dump(yaml_content,f,default_flow_style= False)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T12:10:08.358633Z","iopub.execute_input":"2025-05-07T12:10:08.358916Z","iopub.status.idle":"2025-05-07T12:13:34.347075Z","shell.execute_reply.started":"2025-05-07T12:10:08.358867Z","shell.execute_reply":"2025-05-07T12:13:34.346317Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Build and Train YOLO ","metadata":{}},{"cell_type":"code","source":"from ultralytics import YOLO\nfrom matplotlib.patches import Rectangle\nimport random\nfrom pathlib import Path\n\n\n\nyaml_path = os.path.join(yolo_data_path,'dataset.yaml')\nyolo_weights_dir = \"/kaggle/working/yolo_weights\"\n#yolo_pretrained_weights = \"yolov8n.pt\"\nyolo_pretrained_weights=Path('/kaggle/input/yolo-model2/yolov8n.pt')\nos.makedirs(yolo_weights_dir,exist_ok=True)\n\n\n\nclass YOLOTraining:\n    def __init__(self,yaml_path,pretrained_weight_path):\n        self.yaml_path =yaml_path\n        self.pretrained_weight_path=pretrained_weight_path\n\n    def train(self):\n        model =YOLO(self.pretrained_weight_path)\n        \n        results = model.train(data=self.yaml_path,epochs=30,batch=16,imgsz=640,\n        project=yolo_weights_dir,name='motor_detector',exist_ok=True,\n        patience=5,              # Early stopping if no improvement for 5 epochs\n        save_period=5,           # Save checkpoints every 5 epochs\n        val=True,                # Ensure validation is performed\n        verbose=True )            # Show detailed output during training\n\n        run_dir = os.path.join(yolo_weights_dir,'motor_detector')\n        print('saved weight path',run_dir)\n        return model,results\n\n  \n    def predict(self,model,num_sample=4):\n        \n        data_path = os.path.join(yolo_data_path,yolo_val_img)\n        if not os.path.exists(data_path):\n            print('No data is exist for validation')\n            return\n        val_images=os.listdir(data_path)\n\n        num_samples= min(num_sample,len(val_images))\n        sample_img = random.sample(val_images,num_samples)\n\n        fig,axes=plt.subplots(2,2,figsize=(12,12))\n        ax=axes.flatten()\n        \n        for i , img_file in enumerate(sample_img):\n            img_path =os.path.join(data_path,img_file)\n            results = model.predict(img_path, conf=0.1)[0]\n            img=Image.open(img_path)\n            ax[i].imshow(np.array(img),cmap='gray')\n\n        \n            # draw the actual label rectangle\n            parts= img_file.split('_')\n            x_part = [p for p in parts if p.startswith('x')]\n            y_part = [p for p in parts if p.startswith('y')]\n            x_c= int(x_part[0][1:].split('.')[0])\n            y_c = int(y_part[0][1:])\n            box=24\n            rect_actual = Rectangle((x_c - box//2 , y_c - box//2),box,box,linewidth=1, edgecolor='g', facecolor='none')\n            ax[i].add_patch(rect_actual)\n\n            # predicted label rectangle\n\n            if len(results.boxes) > 0:\n                print('here predicting the result')\n                boxes = results.boxes.xyxy.cpu().numpy()\n                confs = results.boxes.conf.cpu().numpy()\n                \n                for box, conf in zip(boxes, confs):\n                    x1, y1, x2, y2 = box\n                    print('boxes coordinate',[ x1, y1, x2, y2])\n                    rect_pred = Rectangle((x1, y1), x2-x1, y2-y1, \n                                         linewidth=1, edgecolor='r', facecolor='none')\n                    ax[i].add_patch(rect_pred)\n                    ax[i].text(x1, y1-5, f'{conf:.2f}', color='red')\n            ax[i].set_title(f\"Image: {img_file}\\nGround Truth (green) vs Prediction (red)\")\n    \n        plt.tight_layout()\n        \n        # Save the predictions plot\n        plt.savefig(os.path.join('/kaggle/working', 'predictions.png'))\n        plt.show()\n\n#if yolo_pretrained_weights.exists():\nyolo = YOLOTraining(yaml_path,yolo_pretrained_weights)\nmodel,results = yolo.train()\nyolo.predict(model,4)\n# else:\n#     print('model path is not exist')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T12:13:34.349023Z","iopub.execute_input":"2025-05-07T12:13:34.349278Z","iopub.status.idle":"2025-05-07T12:36:26.837926Z","shell.execute_reply.started":"2025-05-07T12:13:34.349260Z","shell.execute_reply":"2025-05-07T12:36:26.836714Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def preload_image_batch(slice_paths):\n        images=[]\n        for path in slice_paths:\n            img=cv2.imread(path)\n            if img is None:\n                img= np.array(Image.open(path))\n            images.append(img)\n        return images\n\nclass GPUProfile:\n    def __init__(self,name):\n        self.name = name\n        self.start_time = None\n    def __enter__(self):\n        if torch.cuda.is_available():\n            torch.cuda.synchronize()\n        self.start_time = time.time()\n        return self\n    def __exit__(self,*args):\n        if torch.cuda.is_available():\n            torch.cuda.synchronize()\n        elapsed= time.time()-self.start_time\n        print(f\"[profile] {self.name}:{elapsed:.3f}s\")\n\n\ndef process_tomogram(tomo_id,model,index=0, total=1):\n    tomo_dir = os.path.join(test_dir,tomo_id)\n    slices_file = sorted([f for f in os.listdir(tomo_dir) if f.endswith('.jpg')])\n    CONCENTRATION=1\n    selected_indices = np.linspace(0,len(slices_file)-1,int(len(slices_file))*CONCENTRATION)\n    selected_indices= np.round(selected_indices).astype(int)\n    slice_files=[slices_file[i] for i in selected_indices]\n    all_detections=[]\n    if device.startswith('cuda'):\n        streams = [torch.cuda.Stream() for _ in range(min(4,BATCH_SIZE))]\n    else:\n        streams=[None]\n    \n    next_batch_thread=None\n    next_batch_image= None\n    \n    for batch_start in range(0,len(slice_files),BATCH_SIZE):\n        if next_batch_thread is not None:\n            next_batch_thread.join()\n            next_batch_image= None\n    \n        batch_end = min(batch_start+ BATCH_SIZE ,len(slice_files))\n        current_slices = slice_files[batch_start:batch_end]\n    \n        next_batch_start = batch_end\n        next_batch_end = min(next_batch_start + BATCH_SIZE, len(slice_files))\n        next_slices = slice_files[next_batch_start:next_batch_end] if next_batch_start < len(slice_files) else []\n    \n        \n    \n        if next_slices:\n            slice_path = [ os.path.join(tomo_dir,f) for f in next_slices]\n            next_batch_thread = threading.Thread(target = preload_image_batch,args=(slice_path,))\n            next_batch_thread.start()\n        else:\n            next_batch_thread=None\n        sub_batches = np.array_split(current_slices,len(streams))\n        sub_batches_result=[]\n        for i , sub_batch in enumerate(sub_batches):\n            if len(sub_batch) == 0:\n                continue\n            stream = streams[i%len(streams)]\n            with torch.cuda.stream(stream) if stream and device.startswith('cuda') else nullcontext:\n                sub_batch_path =[os.path.join(tomo_dir,f) for f in sub_batch if f.endswith('.jpg')]\n                sub_batch_slice_num = [int(slice_file.split('_')[1].split('.')[0]) for slice_file in sub_batch]\n                with GPUProfile(f\"Inference batch{i+1}/{len(sub_batches)}\"):\n                    sub_results = model(sub_batch_path,verbose=False)\n                for j , result in enumerate(sub_results):\n                    if (len(result.boxes))>0:\n                        boxes= result.boxes\n                        for box_id , confidence in enumerate(boxes.conf):\n                            if confidence >= CONFIDENCE_THRESHOLD:\n                                    # Get bounding box coordinates\n                                    x1, y1, x2, y2 = boxes.xyxy[box_idx].cpu().numpy()\n                                    \n                                    # Calculate center coordinates\n                                    x_center = (x1 + x2) / 2\n                                    y_center = (y1 + y2) / 2\n                                    \n                                    # Store detection with 3D coordinates\n                                    all_detections.append({\n                                        'z': round(sub_batch_slice_nums[j]),\n                                        'y': round(y_center),\n                                        'x': round(x_center),\n                                        'confidence': float(confidence)\n                                    })\n        # Synchronize streams\n        if device.startswith('cuda'):\n            torch.cuda.synchronize()\n        \n    # Clean up thread if still running\n    if next_batch_thread is not None:\n        next_batch_thread.join()\n    \n    # 3D Non-Maximum Suppression to merge nearby detections across slices\n    final_detections = perform_3d_nms(all_detections, NMS_IOU_THRESHOLD)\n    \n    # Sort detections by confidence (highest first)\n    final_detections.sort(key=lambda x: x['confidence'], reverse=True)\n    \n    # If there are no detections, return NA values\n    if not final_detections:\n        return {\n            'tomo_id': tomo_id,\n            'Motor axis 0': -1,\n            'Motor axis 1': -1,\n            'Motor axis 2': -1\n        }\n    \n    # Take the detection with highest confidence\n    best_detection = final_detections[0]\n    \n    # Return result with integer coordinates\n    return {\n        'tomo_id': tomo_id,\n        'Motor axis 0': round(best_detection['z']),\n        'Motor axis 1': round(best_detection['y']),\n        'Motor axis 2': round(best_detection['x'])\n    }\n                                \n            \ndef perform_3d_nms(detections, iou_threshold):\n    \"\"\"\n    Perform 3D Non-Maximum Suppression on detections to merge nearby motors\n    \"\"\"\n    if not detections:\n        return []\n    \n    # Sort by confidence (highest first)\n    detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)\n    \n    # List to store final detections after NMS\n    final_detections = []\n    \n    # Define 3D distance function\n    def distance_3d(d1, d2):\n        return np.sqrt((d1['z'] - d2['z'])**2 + \n                       (d1['y'] - d2['y'])**2 + \n                       (d1['x'] - d2['x'])**2)\n    \n    # Maximum distance threshold (based on box size and slice gap)\n    box_size = 24  # Same as annotation box size\n    distance_threshold = box_size * iou_threshold\n    \n    # Process each detection\n    while detections:\n        # Take the detection with highest confidence\n        best_detection = detections.pop(0)\n        final_detections.append(best_detection)\n        \n        # Filter out detections that are too close to the best detection\n        detections = [d for d in detections if distance_3d(d, best_detection) > distance_threshold]\n    \n    return final_detections\n    \n\n    \n\n\n\n\nimport os\nimport pandas as pd\nimport numpy as np\nimport cv2\nfrom PIL import Image\nfrom ultralytics import YOLO\nimport threading\nimport time\nfrom contextlib import nullcontext\nfrom concurrent.futures import ThreadPoolExecutor\nimport torch\nrandom.seed(42)\nnp.random.seed(42)\ntorch.manual_seed(42)\n\ndata_path ='/kaggle/input/byu-locating-bacterial-flagellar-motors-2025'\ntest_dir=os.path.join(data_path,'test')\nmodel_path='/kaggle/working/yolo_weights/motor_detector/weights/best.pt'\n\ndevice = 'cuda:0' if torch.cuda.is_available() else 'cpu'\nBATCH_SIZE = 8\nCONFIDENCE_THRESHOLD = 0.45  # Lower threshold to catch more potential motors\nMAX_DETECTIONS_PER_TOMO = 3  # Keep track of top N detections per tomogram\nNMS_IOU_THRESHOLD = 0.2\n\nif device.startswith('cuda'):\n    torch.backends.cudnn.banchmark = True\n    torch.backends.cudnn.deterministics= False\n    torch.backends.cuda.matmul.allow_tf32 = True\n    torch.backends.cudnn.allow_tf32=True\n\n    gpu_name = torch.cuda.get_device_name(0)\n    gpu_memory = torch.cuda.get_device_properties(0).total_memory/1e9\n    memory_allocated = torch.cuda.memory_allocated(0)/1e9\n    free_memory = gpu_memory - memory_allocated\n    BATCH_SIZE = max(8, min(32, int(free_memory * 4)))\nelse:\n    BATCH_SIZE = 4\n\ndef generate_submission():\n\n    test_tomos = sorted([d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))])\n    total_tomos = len(test_tomos)\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n        \n    model = YOLO(model_path)\n    model.to(device)\n    if device.startswith('cuda'):\n        model.fuse() # run all the layer like conv , batchnorm , maxpool in one time\n        if (torch.cuda.get_device_capability(0)[0]>7) :\n            model=model.half()\n            print('half precision')\n    \n    results = []\n    motor_found = 0\n    \n    with ThreadPoolExecutor(max_workers=1) as executer:\n        future_to_tomo={}\n        for i, tomo_id in enumerate(test_tomos,1):\n            future= executer.submit(process_tomogram,tomo_id,model,i,total_tomos)\n            future_to_tomo[future]=tomo_id\n             # Process completed futures as they complete\n        for future in future_to_tomo:\n            tomo_id = future_to_tomo[future]\n            try:\n                # Clear CUDA cache between tomograms\n                if torch.cuda.is_available():\n                    torch.cuda.empty_cache()\n                    \n                result = future.result()\n                results.append(result)\n                \n                # Update motors found count\n                has_motor = not pd.isna(result['Motor axis 0'])\n                if has_motor:\n                    motors_found += 1\n                    print(f\"Motor found in {tomo_id} at position: \"\n                          f\"z={result['Motor axis 0']}, y={result['Motor axis 1']}, x={result['Motor axis 2']}\")\n                else:\n                    print(f\"No motor detected in {tomo_id}\")\n                    \n                print(f\"Current detection rate: {motors_found}/{len(results)} ({motors_found/len(results)*100:.1f}%)\")\n            \n            except Exception as e:\n                print(f\"Error processing {tomo_id}: {e}\")\n                # Create a default entry for failed tomograms\n            results.append({\n                'tomo_id': tomo_id,\n                'Motor axis 0': -1,\n                'Motor axis 1': -1,\n                'Motor axis 2': -1\n            })\n    \n    # Create submission dataframe\n    submission_df = pd.DataFrame(results)\n    \n    # Ensure proper column order\n    submission_df = submission_df[['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2']]\n    submission_path ='/kaggle/working/sample_submission.csv'\n    \n    # Save the submission file\n    submission_df.to_csv(submission_path, index=False)\n    return submission_df\n\n\n\n# Run the submission pipeline\nif __name__ == \"__main__\":\n    # Time entire process  \n    start_time = time.time()\n    \n    # Generate submission\n    submission = generate_submission()\n    \n    # Print total execution time\n    elapsed = time.time() - start_time\n    print(f\"\\nTotal execution time: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T13:26:32.736713Z","iopub.execute_input":"2025-05-07T13:26:32.737334Z","iopub.status.idle":"2025-05-07T13:27:02.937302Z","shell.execute_reply.started":"2025-05-07T13:26:32.737301Z","shell.execute_reply":"2025-05-07T13:27:02.936534Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}