{"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":103103,"databundleVersionId":13042974,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import glob\nimport argparse\nimport shutil\n\n\nif __name__ == '__main__':\n    inp_folder = '/kaggle/input/alpha-dent/AlphaDent' + '/'\n    out_folder = '/kaggle/output/alpha-dent/AlphaDent_4_classes' + '/'\n\n    shutil.copytree(inp_folder, out_folder)\n\n    # Replace txt files\n    txt_paths = glob.glob(out_folder + '**/*.txt', recursive=True)\n    for txt_path in txt_paths:\n        lines = open(txt_path).readlines()\n        out = open(txt_path, 'w')\n        for line in lines:\n            if line[0] == '4' or line[0] == '5' or line[0] == '6' or line[0] == '7' or line[0] == '8':\n                out.write('3' + line[1:])\n            else:\n                out.write(line)\n        out.close()\n\n    id_to_classes = {\n        1: 'Abrasion',\n        2: 'Filling',\n        3: 'Crown',\n        4: 'Caries',\n    }\n\n    # Create .yaml file\n    out = open(out_folder + 'yolo_seg_train.yaml', 'w')\n    out.write('path: {}\\n'.format(out_folder))\n    out.write('train: images/train\\n')\n    out.write('val: images/valid\\n')\n    out.write('names:\\n')\n    out.write('  0: Abrasion\\n')\n    out.write('  1: Filling\\n')\n    out.write('  2: Crown\\n')\n    out.write('  3: Caries\\n')\n    out.close()\n\n\nprint(\"Done\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T14:06:15.093425Z","iopub.execute_input":"2025-09-12T14:06:15.093796Z","iopub.status.idle":"2025-09-12T14:06:15.141813Z","shell.execute_reply.started":"2025-09-12T14:06:15.093749Z","shell.execute_reply":"2025-09-12T14:06:15.140645Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport sys\nimport time\nimport glob\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nimport yaml\nimport random\nfrom PIL import Image\nimport warnings","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T14:06:21.664023Z","iopub.execute_input":"2025-09-12T14:06:21.664740Z","iopub.status.idle":"2025-09-12T14:06:21.669801Z","shell.execute_reply.started":"2025-09-12T14:06:21.664712Z","shell.execute_reply":"2025-09-12T14:06:21.668837Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define original paths\nBASE_PATH = '/kaggle/output/alpha-dent/AlphaDent_4_classes'\nORIGINAL_TRAIN_IMAGES_PATH = f'{BASE_PATH}/images/train'\nVALID_IMAGES_PATH = f'{BASE_PATH}/images/valid'\nTEST_IMAGES_PATH = f'{BASE_PATH}/images/test'\nORIGINAL_TRAIN_LABELS_PATH = f'{BASE_PATH}/labels/train'\nVALID_LABELS_PATH = f'{BASE_PATH}/labels/valid'\n\n# Output paths\nOUTPUT_DIR = '/kaggle/working/'\nWEIGHTS_DIR = f'{OUTPUT_DIR}/weights'\nos.makedirs(WEIGHTS_DIR, exist_ok=True)\n\n# Create new dataset structure with 90/10 split\nNEW_DATASET_PATH = f'{OUTPUT_DIR}/alphadent_90_10_split'\nNEW_TRAIN_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/train'\nNEW_EVAL_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/eval'\nNEW_VALID_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/valid'\nNEW_TEST_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/test'\nNEW_TRAIN_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/train'\nNEW_EVAL_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/eval'\nNEW_VALID_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/valid'\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T14:06:24.963058Z","iopub.execute_input":"2025-09-12T14:06:24.963349Z","iopub.status.idle":"2025-09-12T14:06:24.969883Z","shell.execute_reply.started":"2025-09-12T14:06:24.963326Z","shell.execute_reply":"2025-09-12T14:06:24.968820Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport sys\nimport time\nimport glob\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nimport yaml\nimport random\nfrom PIL import Image\nimport shutil\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Disable wandb\nos.environ['WANDB_DISABLED'] = 'true'\n\n# Set random seeds for reproducibility\nrandom.seed(42)\nnp.random.seed(42)\n\n# Install required packages\nprint(\"Installing required packages...\")\nos.system('pip install -q ultralytics')\n\nimport torch\nfrom ultralytics import YOLO\n\n# Set deterministic behavior for PyTorch\ntorch.manual_seed(42)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(42)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nprint(f'\\nPyTorch Version: {torch.__version__}')\nprint(f'CUDA Available: {torch.cuda.is_available()}')\nif torch.cuda.is_available():\n    print(f'CUDA Device: {torch.cuda.get_device_name(0)}')\n\n# Define original paths\nBASE_PATH = '/kaggle/output/alpha-dent/AlphaDent_4_classes'\nORIGINAL_TRAIN_IMAGES_PATH = f'{BASE_PATH}/images/train'\nVALID_IMAGES_PATH = f'{BASE_PATH}/images/valid'\nTEST_IMAGES_PATH = f'{BASE_PATH}/images/test'\nORIGINAL_TRAIN_LABELS_PATH = f'{BASE_PATH}/labels/train'\nVALID_LABELS_PATH = f'{BASE_PATH}/labels/valid'\n\n# Output paths\nOUTPUT_DIR = '/kaggle/working/'\nWEIGHTS_DIR = f'{OUTPUT_DIR}/weights'\nos.makedirs(WEIGHTS_DIR, exist_ok=True)\n\n# Create new dataset structure with 90/10 split\nNEW_DATASET_PATH = f'{OUTPUT_DIR}/alphadent_90_10_split'\nNEW_TRAIN_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/train'\nNEW_EVAL_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/eval'\nNEW_VALID_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/valid'\nNEW_TEST_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/test'\nNEW_TRAIN_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/train'\nNEW_EVAL_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/eval'\nNEW_VALID_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/valid'\n\n# Create directories\nos.makedirs(NEW_TRAIN_IMAGES_PATH, exist_ok=True)\nos.makedirs(NEW_EVAL_IMAGES_PATH, exist_ok=True)\nos.makedirs(NEW_VALID_IMAGES_PATH, exist_ok=True)\nos.makedirs(NEW_TEST_IMAGES_PATH, exist_ok=True)\nos.makedirs(NEW_TRAIN_LABELS_PATH, exist_ok=True)\nos.makedirs(NEW_EVAL_LABELS_PATH, exist_ok=True)\nos.makedirs(NEW_VALID_LABELS_PATH, exist_ok=True)\n\nprint(\"\\n=== Creating 90/10 Train/Eval Split ===\")\n\n# Get all training images\noriginal_train_images = sorted(glob.glob(f'{ORIGINAL_TRAIN_IMAGES_PATH}/*.jpg'))\nprint(f\"Total original training images: {len(original_train_images)}\")\n\n# Shuffle and split into 90% train, 10% eval\nrandom.shuffle(original_train_images)\nsplit_idx = int(0.9 * len(original_train_images))\ntrain_90_images = original_train_images[:split_idx]\neval_10_images = original_train_images[split_idx:]\n\nprint(f\"90% for training: {len(train_90_images)}\")\nprint(f\"10% for evaluation: {len(eval_10_images)}\")\n\ndef copy_files(image_list, dest_images_dir, dest_labels_dir, source_labels_dir, desc):\n    \"\"\"Copy images and corresponding labels to destination directories.\"\"\"\n    for img_path in tqdm(image_list, desc=desc):\n        # Copy image\n        img_filename = os.path.basename(img_path)\n        shutil.copy2(img_path, os.path.join(dest_images_dir, img_filename))\n        \n        # Copy corresponding label\n        label_filename = img_filename.replace('.jpg', '.txt')\n        source_label_path = os.path.join(source_labels_dir, label_filename)\n        dest_label_path = os.path.join(dest_labels_dir, label_filename)\n        \n        if os.path.exists(source_label_path):\n            shutil.copy2(source_label_path, dest_label_path)\n\n# Copy 90% training data\ncopy_files(train_90_images, NEW_TRAIN_IMAGES_PATH, NEW_TRAIN_LABELS_PATH, \n          ORIGINAL_TRAIN_LABELS_PATH, \"Copying 90% training data\")\n\n# Copy 10% evaluation data\ncopy_files(eval_10_images, NEW_EVAL_IMAGES_PATH, NEW_EVAL_LABELS_PATH, \n          ORIGINAL_TRAIN_LABELS_PATH, \"Copying 10% evaluation data\")\n\n# Copy original validation data (unchanged)\noriginal_valid_images = glob.glob(f'{VALID_IMAGES_PATH}/*.jpg')\ncopy_files(original_valid_images, NEW_VALID_IMAGES_PATH, NEW_VALID_LABELS_PATH, \n          VALID_LABELS_PATH, \"Copying validation data\")\n\n# Copy test data (images only, no labels)\nprint(\"Copying test images...\")\ntest_images = glob.glob(f'{TEST_IMAGES_PATH}/*.jpg')\nfor img_path in tqdm(test_images, desc=\"Copying test images\"):\n    img_filename = os.path.basename(img_path)\n    shutil.copy2(img_path, os.path.join(NEW_TEST_IMAGES_PATH, img_filename))\n\n# Define class information\nCLASS_INFO = {\n    0: {'name': 'Abrasion', 'description': 'Teeth with mechanical wear of hard tissues'},\n    1: {'name': 'Filling', 'description': 'Dental fillings of various types'},\n    2: {'name': 'Crown', 'description': 'Dental crown (restoration)'},\n    3: {'name': 'Caries', 'description': 'Caries in fissures and pits'}\n}\n\n# Create YAML configuration for YOLO\nprint(\"\\n=== Creating YOLO Configuration ===\")\nyolo_config = {\n    'path': NEW_DATASET_PATH,\n    'train': 'images/train',\n    'val': 'images/valid',\n    'test': 'images/test',\n    'nc': 4,\n    'names': [CLASS_INFO[i]['name'] for i in range(4)]\n}\n\n# Save the configuration\nCUSTOM_YAML_PATH = f'{OUTPUT_DIR}/alphadent_config_90_10.yaml'\nwith open(CUSTOM_YAML_PATH, 'w') as f:\n    yaml.dump(yolo_config, f, default_flow_style=False)\nprint(f\"Created custom YAML config at: {CUSTOM_YAML_PATH}\")\n\n# Count images in new dataset\nnew_train_images = sorted(glob.glob(f'{NEW_TRAIN_IMAGES_PATH}/*.jpg'))\nnew_eval_images = sorted(glob.glob(f'{NEW_EVAL_IMAGES_PATH}/*.jpg'))\nnew_valid_images = sorted(glob.glob(f'{NEW_VALID_IMAGES_PATH}/*.jpg'))\nnew_test_images = sorted(glob.glob(f'{NEW_TEST_IMAGES_PATH}/*.jpg'))\n\nprint(f\"\\n=== New Dataset Statistics ===\")\nprint(f\"Training images (90%): {len(new_train_images)}\")\nprint(f\"Evaluation images (10%): {len(new_eval_images)}\")\nprint(f\"Validation images: {len(new_valid_images)}\")\nprint(f\"Test images: {len(new_test_images)}\")\n\n# Analyze class distribution\ndef analyze_class_distribution(labels_path, dataset_name):\n    \"\"\"Analyze class distribution in dataset.\"\"\"\n    class_counts = {i: 0 for i in range(4)}\n    total_annotations = 0\n    \n    label_files = glob.glob(f'{labels_path}/*.txt')\n    \n    for label_file in tqdm(label_files, desc=f\"Analyzing {dataset_name} labels\", leave=False):\n        if os.path.exists(label_file) and os.path.getsize(label_file) > 0:\n            try:\n                with open(label_file, 'r') as f:\n                    lines = f.readlines()\n                    for line in lines:\n                        if line.strip():\n                            parts = line.strip().split()\n                            if parts:\n                                class_id = int(parts[0])\n                                if 0 <= class_id < 4:\n                                    class_counts[class_id] += 1\n                                    total_annotations += 1\n            except Exception:\n                continue\n    \n    return class_counts, total_annotations\n\nprint(\"\\n=== Analyzing Class Distribution ===\")\ntrain_class_counts, train_total = analyze_class_distribution(NEW_TRAIN_LABELS_PATH, \"training\")\neval_class_counts, eval_total = analyze_class_distribution(NEW_EVAL_LABELS_PATH, \"evaluation\")\nvalid_class_counts, valid_total = analyze_class_distribution(NEW_VALID_LABELS_PATH, \"validation\")\n\nprint(f\"Training set (90%): {train_total} total annotations\")\nfor i, (class_id, count) in enumerate(train_class_counts.items()):\n    print(f\"  {CLASS_INFO[class_id]['name']}: {count}\")\n\nprint(f\"Evaluation set (10%): {eval_total} total annotations\")\nfor i, (class_id, count) in enumerate(eval_class_counts.items()):\n    print(f\"  {CLASS_INFO[class_id]['name']}: {count}\")\n\nprint(f\"Validation set: {valid_total} total annotations\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T14:14:45.125038Z","iopub.execute_input":"2025-09-12T14:14:45.125414Z","iopub.status.idle":"2025-09-12T14:15:28.064056Z","shell.execute_reply.started":"2025-09-12T14:14:45.125363Z","shell.execute_reply":"2025-09-12T14:15:28.062909Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Training configuration\nprint(\"\\n=== Model Training Configuration ===\")\nEPOCHS = 30\nIMAGE_SIZE = 640\nBATCH_SIZE = 8 if torch.cuda.is_available() else 4\nPATIENCE = 5\n\nprint(f\"Epochs: {EPOCHS}\")\nprint(f\"Image Size: {IMAGE_SIZE}\")\nprint(f\"Batch Size: {BATCH_SIZE}\")\nprint(f\"Early Stopping Patience: {PATIENCE}\")\n\n# Initialize and train model\nprint(\"\\n=== Starting Model Training ===\")\nmodel = YOLO('yolov8x-seg.pt')\n\n# Train the model with optimized parameters\nresults = model.train(\n    data=CUSTOM_YAML_PATH,\n    epochs=EPOCHS,\n    imgsz=IMAGE_SIZE,\n    batch=BATCH_SIZE,\n    patience=PATIENCE,\n    save=True,\n    save_period=10,\n    project=OUTPUT_DIR,\n    name='alphadent_yolov8x_90_10',\n    exist_ok=True,\n    pretrained=True,\n    optimizer='AdamW',\n    lr0=0.001,\n    lrf=0.01,\n    momentum=0.937,\n    weight_decay=0.0005,\n    warmup_epochs=3.0,\n    warmup_momentum=0.8,\n    warmup_bias_lr=0.1,\n    box=7.5,\n    cls=0.5,\n    dfl=1.5,\n    hsv_h=0.015,\n    hsv_s=0.7,\n    hsv_v=0.4,\n    degrees=0.0,\n    translate=0.1,\n    scale=0.5,\n    shear=0.0,\n    perspective=0.0,\n    flipud=0.0,\n    fliplr=0.5,\n    mosaic=1.0,\n    mixup=0.0,\n    copy_paste=0.0,\n    plots=True,\n    device=0 if torch.cuda.is_available() else 'cpu',\n    workers=2,\n    verbose=True,\n    amp=True,\n    val=True\n)\n\nprint(\"\\nTraining completed!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T14:15:58.888826Z","iopub.execute_input":"2025-09-12T14:15:58.889957Z","iopub.status.idle":"2025-09-12T14:16:11.251135Z","shell.execute_reply.started":"2025-09-12T14:15:58.889920Z","shell.execute_reply":"2025-09-12T14:16:11.249446Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load best model\nprint(\"=== Loading Best Model ===\")\nbest_model_path = f'{OUTPUT_DIR}/alphadent_yolov8x_90_10/weights/best.pt'\nif os.path.exists(best_model_path):\n    model = YOLO(best_model_path)\n    print(f\"Loaded best model from: {best_model_path}\")\nelse:\n    last_model_path = f'{OUTPUT_DIR}/alphadent_yolov8x_90_10/weights/last.pt'\n    if os.path.exists(last_model_path):\n        model = YOLO(last_model_path)\n        print(f\"Loaded last model from: {last_model_path}\")\n    else:\n        print(\"Warning: No trained model found, using pretrained model\")\n        model = YOLO('yolov8x-seg.pt')\n\n# Validate model on original validation set\nprint(\"\\n=== Model Validation on Original Validation Set ===\")\ntry:\n    metrics = model.val(\n        data=CUSTOM_YAML_PATH,\n        imgsz=IMAGE_SIZE,\n        batch=1,\n        conf=0.001,\n        iou=0.5,\n        max_det=300,\n        device=0 if torch.cuda.is_available() else 'cpu',\n        plots=False,\n        save_json=False,\n    )\n    \n    print(f\"\\nValidation Results on Original Validation Set:\")\n    print(f\"mAP@50: {metrics.seg.map50:.4f}\")\n    print(f\"mAP@50-95: {metrics.seg.map:.4f}\")\nexcept Exception as e:\n    print(f\"Validation error (non-critical): {e}\")\n\n# Create custom YAML for evaluation on the 10% held-out data\neval_config = {\n    'path': NEW_DATASET_PATH,\n    'train': 'images/train',\n    'val': 'images/eval',  # Point to eval set for validation\n    'test': 'images/test',\n    'nc': 4,\n    'names': [CLASS_INFO[i]['name'] for i in range(4)]\n}\n\nEVAL_YAML_PATH = f'{OUTPUT_DIR}/alphadent_eval_config.yaml'\nwith open(EVAL_YAML_PATH, 'w') as f:\n    yaml.dump(eval_config, f, default_flow_style=False)\n\n# Evaluate model on the 10% held-out labeled data\nprint(\"\\n=== Model Evaluation on 10% Held-out Labeled Data ===\")\ntry:\n    eval_metrics = model.val(\n        data=EVAL_YAML_PATH,\n        imgsz=IMAGE_SIZE,\n        batch=1,\n        conf=0.001,\n        iou=0.5,\n        max_det=300,\n        device=0 if torch.cuda.is_available() else 'cpu',\n        plots=True,\n        save_json=True,\n        name='eval_10_percent'\n    )\n    \n    print(f\"\\nEvaluation Results on 10% Held-out Data:\")\n    print(f\"mAP@50: {eval_metrics.seg.map50:.4f}\")\n    print(f\"mAP@50-95: {eval_metrics.seg.map:.4f}\")\n    print(f\"Per-class mAP@50:\")\n    for i, map_val in enumerate(eval_metrics.seg.maps):\n        print(f\"  {CLASS_INFO[i]['name']}: {map_val:.4f}\")\n        \nexcept Exception as e:\n    print(f\"Evaluation error: {e}\")\n\n# Final summary\nprint(\"\\n=== Summary ===\")\nprint(f\"✓ Successfully created 90/10 train/eval split\")\nprint(f\"✓ Trained model on 90% of original training data ({len(new_train_images)} images)\")\nprint(f\"✓ Evaluated model on 10% held-out labeled data ({len(new_eval_images)} images)\")\nprint(f\"✓ Also validated on original validation set ({len(new_valid_images)} images)\")\nprint(f\"✓ Model weights saved to: {OUTPUT_DIR}/alphadent_yolov8x_90_10/weights/\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T14:16:14.178474Z","iopub.execute_input":"2025-09-12T14:16:14.178816Z","iopub.status.idle":"2025-09-12T14:55:26.619695Z","shell.execute_reply.started":"2025-09-12T14:16:14.178790Z","shell.execute_reply":"2025-09-12T14:55:26.618434Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =====================\n# Train Faster R-CNN\n# =====================\nprint(\"=== Training Faster R-CNN ===\")\n\nimport glob, os\nimport torch\nimport torchvision\nimport torch.optim as optim\nimport torchvision.transforms as T\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nfrom tqdm import tqdm\n\n# Custom dataset class for YOLO-format labels\nclass YOLODataset(Dataset):\n    def __init__(self, images_dir, labels_dir, transforms=None):\n        self.images = sorted(glob.glob(f\"{images_dir}/*.jpg\"))\n        self.labels_dir = labels_dir\n        self.transforms = transforms\n\n    def __getitem__(self, idx):\n        img_path = self.images[idx]\n        label_path = os.path.join(\n            self.labels_dir, os.path.basename(img_path).replace('.jpg', '.txt')\n        )\n\n        img = Image.open(img_path).convert(\"RGB\")\n        w, h = img.size\n\n        boxes = []\n        labels = []\n\n        if os.path.exists(label_path):\n            with open(label_path, 'r') as f:\n                for line in f.readlines():\n                    parts = list(map(float, line.strip().split()))\n                    cls = int(parts[0])\n                    if len(parts) >= 5:  # YOLO bbox format\n                        x, y, bw, bh = parts[1:5]\n                        x1 = (x - bw/2) * w\n                        y1 = (y - bh/2) * h\n                        x2 = (x + bw/2) * w\n                        y2 = (y + bh/2) * h\n\n                        # clamp to image size\n                        x1, y1 = max(0, x1), max(0, y1)\n                        x2, y2 = min(w, x2), min(h, y2)\n\n                        # only keep valid boxes\n                        if x2 > x1 and y2 > y1:\n                            boxes.append([x1, y1, x2, y2])\n                            labels.append(cls + 1)\n\n        target = {\n            \"boxes\": torch.tensor(boxes, dtype=torch.float32),\n            \"labels\": torch.tensor(labels, dtype=torch.int64)\n        }\n\n        if self.transforms:\n            img = self.transforms(img)\n\n        return img, target\n\n    def __len__(self):\n        return len(self.images)\n\n\n\n# Define transforms\ntransform = T.Compose([T.ToTensor()])\n\n# Load datasets\ntrain_dataset = YOLODataset(\n    f\"{OUTPUT_DIR}/alphadent_90_10_split/images/train\",\n    f\"{OUTPUT_DIR}/alphadent_90_10_split/labels/train\",\n    transforms=transform\n)\nval_dataset = YOLODataset(\n    f\"{OUTPUT_DIR}/alphadent_90_10_split/images/valid\",\n    f\"{OUTPUT_DIR}/alphadent_90_10_split/labels/valid\",\n    transforms=transform\n)\n\ntrain_loader = DataLoader(\n    train_dataset, batch_size=8, shuffle=True,\n    collate_fn=lambda x: tuple(zip(*x))\n)\nval_loader = DataLoader(\n    val_dataset, batch_size=8, shuffle=False,\n    collate_fn=lambda x: tuple(zip(*x))\n)\n\n# Load Faster R-CNN\nnum_classes = 5  # 4 classes + background\nmodel_rcnn = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=\"DEFAULT\")\nin_features = model_rcnn.roi_heads.box_predictor.cls_score.in_features\nmodel_rcnn.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)\n\n# Optimizer\nparams = [p for p in model_rcnn.parameters() if p.requires_grad]\noptimizer = optim.AdamW(params, lr=0.001)\nimport copy\n\n# Training loop with validation and model checkpointing\nnum_epochs = 5\ndevice = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')\nmodel_rcnn.to(device)\n\nbest_loss = float(\"inf\")\nbest_model_wts = copy.deepcopy(model_rcnn.state_dict())\n\nfor epoch in range(num_epochs):\n    # --------------------\n    # Training\n    # --------------------\n    model_rcnn.train()\n    total_loss = 0\n    for images, targets in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs} [Train]\", leave=False):\n        images = list(img.to(device) for img in images)\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n        loss_dict = model_rcnn(images, targets)\n        losses = sum(loss for loss in loss_dict.values())\n\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n        total_loss += losses.item()\n\n    avg_train_loss = total_loss / len(train_loader)\n    print(f\"Epoch {epoch+1} - Training Loss: {avg_train_loss:.4f}\")\n\n    # --------------------\n    # Validation\n    # --------------------\n    model_rcnn.eval()\n    val_loss = 0\n    for images, targets in tqdm(val_loader, desc=f\"Epoch {epoch+1}/{num_epochs} [Val]\", leave=False):\n        images = list(img.to(device) for img in images)\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n        with torch.no_grad():\n            loss_dict = model_rcnn(images, targets)\n            losses = sum(loss for loss in loss_dict.values())\n            val_loss += losses.item()\n\n    avg_val_loss = val_loss / len(val_loader)\n    print(f\"Epoch {epoch+1} - Validation Loss: {avg_val_loss:.4f}\")\n\n\n    # --------------------\n    # Save Best Model\n    # --------------------\n    if avg_val_loss < best_loss:\n        best_loss = avg_val_loss\n        best_model_wts = copy.deepcopy(model_rcnn.state_dict())\n        torch.save(model_rcnn.state_dict(), \"best_faster_rcnn.pth\")\n        print(f\"✅ Saved new best model at epoch {epoch+1} (val_loss={avg_val_loss:.4f})\")\n\n# Load best weights back into model\nmodel_rcnn.load_state_dict(best_model_wts)\nprint(\"🏆 Best model loaded from training!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T14:14:16.472504Z","iopub.status.idle":"2025-09-12T14:14:16.472810Z","shell.execute_reply.started":"2025-09-12T14:14:16.472668Z","shell.execute_reply":"2025-09-12T14:14:16.472684Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"yolo_model = model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T14:55:43.826708Z","iopub.execute_input":"2025-09-12T14:55:43.827214Z","iopub.status.idle":"2025-09-12T14:55:43.833618Z","shell.execute_reply.started":"2025-09-12T14:55:43.827183Z","shell.execute_reply":"2025-09-12T14:55:43.832310Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torchvision\nimport torchvision.transforms as T\n\ndef load_best_rcnn_model(model_path=\"best_faster_rcnn.pth\", num_classes=5, device=None):\n    \"\"\"\n    Load the best trained Faster R-CNN model\n    \"\"\"\n    if device is None:\n        device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    \n    # Create the same model architecture as in training\n    model_rcnn = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=\"DEFAULT\")\n    \n    # Modify the classifier head to match your training setup\n    in_features = model_rcnn.roi_heads.box_predictor.cls_score.in_features\n    model_rcnn.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(\n        in_features, num_classes\n    )\n    \n    # Load the trained weights\n    model_rcnn.load_state_dict(torch.load(model_path, map_location=device))\n    \n    # Set to evaluation mode\n    model_rcnn.eval()\n    \n    # Move to device\n    model_rcnn = model_rcnn.to(device)\n    \n    print(f\"✅ Loaded best R-CNN model from {model_path}\")\n    return model_rcnn\n\n# Load your trained models\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\n# Load R-CNN model\nrcnn_model = load_best_rcnn_model(\"best_faster_rcnn.pth\", num_classes=5, device=device)\n\n# Load your other models (YOLO and classifier)\nclassifier = torchvision.models.resnet18(pretrained=False)\nclassifier.fc = torch.nn.Linear(classifier.fc.in_features, 6)\n\nclassifier.load_state_dict(torch.load(\"caries_polygon_classifier.pth\", map_location=device))\nclassifier.eval()\n\nyolo_model = yolo_model.to(device)\nclassifier = classifier.to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T15:00:04.800718Z","iopub.execute_input":"2025-09-12T15:00:04.801122Z","iopub.status.idle":"2025-09-12T15:00:06.080953Z","shell.execute_reply.started":"2025-09-12T15:00:04.801097Z","shell.execute_reply":"2025-09-12T15:00:06.080167Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## CLASSIFIER","metadata":{}},{"cell_type":"code","source":"CLASS_INFO = {\n    0: {'name': 'Abrasion', 'description': 'Teeth with mechanical wear of hard tissues'},\n    1: {'name': 'Filling', 'description': 'Dental fillings of various types'},\n    2: {'name': 'Crown', 'description': 'Dental crown (restoration)'},\n    3: {'name': 'Caries 1', 'description': 'Caries in fissures and pits (occlusal surfaces of molars/premolars, buccal surfaces of molars, lingual surfaces of upper incisors).'},\n    4: {'name': 'Caries 2', 'description': 'Caries on proximal surfaces of molars and premolars.'},\n    5: {'name': 'Caries 3', 'description': 'Caries on proximal surfaces of incisors/canines without incisal edge involvement.'},\n    6: {'name': 'Caries 4', 'description': 'Caries on proximal surfaces of incisors/canines with incisal edge involvement.'},\n    7: {'name': 'Caries 5', 'description': 'Cervical caries (buccal/lingual surfaces).'},\n    8: {'name': 'Caries 6', 'description': 'Caries on incisal edges of anterior teeth or cusps of molars.'},\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T15:01:10.243958Z","iopub.execute_input":"2025-09-12T15:01:10.244316Z","iopub.status.idle":"2025-09-12T15:01:10.250409Z","shell.execute_reply.started":"2025-09-12T15:01:10.244288Z","shell.execute_reply":"2025-09-12T15:01:10.249416Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import cv2\nimport torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport os\nfrom torchvision import transforms\nfrom PIL import Image\nimport pandas as pd\n\n# ----- Transform for classifier -----\nclf_transform = transforms.Compose([\n    transforms.Resize((224, 224)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.5]*3, [0.5]*3)\n])\n\ndef run_rcnn_on_crops(rcnn_model, image, boxes, device, conf_threshold=0.5):\n    \"\"\"\n    Run R-CNN model on cropped regions defined by bounding boxes\n    \"\"\"\n    results = []\n    \n    for box in boxes:\n        x1, y1, x2, y2 = box.astype(int)\n        \n        # Crop the image\n        crop = image[y1:y2, x1:x2]\n        \n        if crop.size > 0:\n            # Convert crop to tensor format expected by R-CNN\n            crop_pil = Image.fromarray(crop).convert(\"RGB\")\n            crop_tensor = transforms.ToTensor()(crop_pil).unsqueeze(0).to(device)\n            \n            # Run R-CNN inference\n            with torch.no_grad():\n                rcnn_outputs = rcnn_model(crop_tensor)\n            \n            # Filter predictions by confidence threshold\n            rcnn_result = rcnn_outputs[0]\n            scores = rcnn_result['scores']\n            keep_indices = scores > conf_threshold\n            \n            filtered_result = {\n                'boxes': rcnn_result['boxes'][keep_indices] if 'boxes' in rcnn_result else None,\n                'labels': rcnn_result['labels'][keep_indices] if 'labels' in rcnn_result else None,\n                'scores': rcnn_result['scores'][keep_indices] if 'scores' in rcnn_result else None,\n            }\n            \n            results.append(filtered_result)\n        else:\n            results.append({'boxes': None, 'labels': None, 'scores': None})\n    \n    return results\n\ndef create_mask_from_rcnn_box(box, crop_shape):\n    \"\"\"\n    Create a simple rectangular mask from R-CNN bounding box\n    Since Faster R-CNN doesn't output segmentation masks by default,\n    we'll create a rectangular mask from the bounding box\n    \"\"\"\n    h, w = crop_shape[:2]\n    mask = np.zeros((h, w), dtype=np.uint8)\n    \n    x1, y1, x2, y2 = box.astype(int)\n    # Clamp coordinates to crop boundaries\n    x1, y1 = max(0, x1), max(0, y1)\n    x2, y2 = min(w, x2), min(h, y2)\n    \n    # Fill the rectangular region\n    mask[y1:y2, x1:x2] = 255\n    \n    return mask\n\ndef predict_and_save_pipeline(yolo_model, rcnn_model, classifier, image_paths, class_names, submission_file=\"final_model_submission.csv\", imgsz=640):\n    all_preds = []\n    row_id = 0\n    device = next(rcnn_model.parameters()).device\n\n    # Move classifier to CPU for efficiency\n    classifier_cpu = classifier.to(\"cpu\")\n\n    for img_path in image_paths:\n        print(f\"Processing: {img_path}\")\n        \n        # Step 1: YOLO for bounding box detection\n        yolo_results = yolo_model.predict(\n            img_path, imgsz=imgsz, conf=0.25, iou=0.45,\n            retina_masks=False, verbose=False  # We don't need YOLO masks\n        )\n        yolo_result = yolo_results[0]\n\n        h, w = yolo_result.orig_shape\n        if yolo_result.boxes is None or len(yolo_result.boxes) == 0:\n            continue\n\n        # Get bounding boxes from YOLO\n        boxes = yolo_result.boxes.xyxy.cpu().numpy()\n        yolo_classes = yolo_result.boxes.cls.cpu().numpy().astype(int)\n        confidences = yolo_result.boxes.conf.cpu().numpy()\n\n        # Step 2: Load image for R-CNN processing\n        img_bgr = cv2.imread(img_path)\n        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)\n\n        # Step 3: Run R-CNN on the bounding box regions\n        rcnn_results = run_rcnn_on_crops(rcnn_model, img_rgb, boxes, device)\n\n        # Step 4: Process each detection\n        for i, (box, yolo_cls, conf) in enumerate(zip(boxes, yolo_classes, confidences)):\n            x1, y1, x2, y2 = box.astype(int)\n\n            final_cls = int(yolo_cls)\n\n            # Step 5: Refine classification using your classifier (for specific classes)\n            if yolo_cls == 3:  # or whatever condition you want\n                crop = img_rgb[y1:y2, x1:x2]\n\n                if crop.size > 0:\n                    crop_pil = Image.fromarray(crop).convert(\"RGB\")\n                    crop_tensor = clf_transform(crop_pil).unsqueeze(0)  # keep on CPU\n\n                    with torch.no_grad():\n                        out = classifier_cpu(crop_tensor)\n                        fine_cls = int(out.argmax(1).item())\n                    final_cls = fine_cls + 3  # shift to [3–8]\n\n            # Step 6: Create polygon from R-CNN result\n            poly_str = \"\"\n            if i < len(rcnn_results):\n                rcnn_result = rcnn_results[i]\n                \n                if (rcnn_result['boxes'] is not None and \n                    len(rcnn_result['boxes']) > 0):\n                    \n                    # Use the highest confidence detection from R-CNN\n                    best_idx = torch.argmax(rcnn_result['scores']).item()\n                    best_box = rcnn_result['boxes'][best_idx].cpu().numpy()\n                    \n                    # Create mask from the R-CNN bounding box\n                    crop_shape = img_rgb[y1:y2, x1:x2].shape\n                    mask = create_mask_from_rcnn_box(best_box, crop_shape)\n                    \n                    # Convert mask to polygon\n                    if mask.sum() > 0:\n                        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n                        \n                        if contours:\n                            largest_contour = max(contours, key=cv2.contourArea)\n                            epsilon = 0.002 * cv2.arcLength(largest_contour, True)\n                            approx_contour = cv2.approxPolyDP(largest_contour, epsilon, True)\n                            \n                            # Convert to original image coordinates\n                            polygon_points = approx_contour.reshape(-1, 2)\n                            polygon_points[:, 0] += x1  # Add x offset\n                            polygon_points[:, 1] += y1  # Add y offset\n                            \n                            # Normalize coordinates\n                            if len(polygon_points) >= 3:\n                                normalized_coords = []\n                                for point in polygon_points:\n                                    x_norm = float(point[0]) / w\n                                    y_norm = float(point[1]) / h\n                                    normalized_coords.extend([\n                                        max(0.0, min(1.0, x_norm)),\n                                        max(0.0, min(1.0, y_norm))\n                                    ])\n                                poly_str = \" \".join([f\"{coord:.6f}\" for coord in normalized_coords])\n\n            # If no valid polygon from R-CNN, create one from YOLO box\n            if not poly_str:\n                # Create rectangle polygon from YOLO box\n                x1_norm, y1_norm = x1/w, y1/h\n                x2_norm, y2_norm = x2/w, y2/h\n                normalized_coords = [x1_norm, y1_norm, x2_norm, y1_norm, x2_norm, y2_norm, x1_norm, y2_norm]\n                poly_str = \" \".join([f\"{coord:.6f}\" for coord in normalized_coords])\n\n            patient_id = os.path.splitext(os.path.basename(img_path))[0]\n            all_preds.append({\n                \"id\": row_id,\n                \"patient_id\": patient_id,\n                \"class_id\": final_cls,\n                \"confidence\": float(conf),\n                \"poly\": poly_str\n            })\n            row_id += 1\n\n    df = pd.DataFrame(all_preds)\n    df.to_csv(submission_file, index=False)\n    print(f\"✅ Submission saved to {submission_file}\")\n\n# Usage\npredict_and_save_pipeline(\n    yolo_model=yolo_model,\n    rcnn_model=rcnn_model,\n    classifier=classifier,\n    image_paths=test_images,\n    class_names=[CLASS_INFO[i]['name'] for i in range(9)],\n    submission_file=\"submission.csv\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T15:01:12.675149Z","iopub.execute_input":"2025-09-12T15:01:12.675531Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\ntorch.cuda.empty_cache()\n# Inference on test set\nprint(\"\\n=== Running Inference on Test Set ===\")\n\ndef convert_to_submission_format(results, image_paths):\n    submission_rows = []\n    \n    for idx, result in enumerate(results):\n        # Get image ID (filename without extension)\n        image_id = os.path.basename(image_paths[idx]).replace('.jpg', '')\n        \n        if result.masks is not None and len(result.masks) > 0:\n            try:\n                # Get masks, classes, and confidences\n                masks = result.masks.xy\n                classes = result.boxes.cls.cpu().numpy().astype(int)\n                confidences = result.boxes.conf.cpu().numpy()\n                h, w = result.orig_shape\n                \n                # Process each detection\n                for mask_idx in range(len(masks)):\n                    if mask_idx < len(classes) and mask_idx < len(confidences):\n                        polygon = masks[mask_idx]\n                        \n                        if len(polygon) >= 3:  # Valid polygon (at least 3 points)\n                            # Normalize coordinates to [0, 1]\n                            normalized_coords = []\n                            for point in polygon:\n                                x_norm = float(point[0]) / w\n                                y_norm = float(point[1]) / h\n                                # Ensure coordinates are within [0, 1]\n                                x_norm = max(0.0, min(1.0, x_norm))\n                                y_norm = max(0.0, min(1.0, y_norm))\n                                normalized_coords.extend([x_norm, y_norm])\n                            \n                            # Format polygon string\n                            poly_str = ' '.join([f'{coord:.6f}' for coord in normalized_coords])\n                            \n                            submission_rows.append({\n                                'patient_id': image_id,\n                                'class_id': int(classes[mask_idx]),\n                                'confidence': float(confidences[mask_idx]),\n                                'poly': poly_str\n                            })\n            except Exception as e:\n                print(f\"Error processing result for image {idx}: {e}\")\n                continue\n    \n    return submission_rows\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T16:26:10.483818Z","iopub.execute_input":"2025-09-07T16:26:10.484075Z","iopub.status.idle":"2025-09-07T16:26:10.496348Z","shell.execute_reply.started":"2025-09-07T16:26:10.484052Z","shell.execute_reply":"2025-09-07T16:26:10.495596Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Process test images\ntest_images = sorted(glob.glob(f'{TEST_IMAGES_PATH}/*.jpg'))\nall_submission_rows = []\nINFERENCE_BATCH_SIZE = 2 if torch.cuda.is_available() else 4\n\nprint(f\"Processing {len(test_images)} test images...\")\n\n# Process in batches\nfor i in tqdm(range(0, len(test_images), INFERENCE_BATCH_SIZE)):\n    torch.cuda.empty_cache()\n    batch_images = test_images[i:i + INFERENCE_BATCH_SIZE]\n    \n    try:\n        # Run inference\n        results = model.predict(\n            batch_images,\n            imgsz=IMAGE_SIZE,\n            conf=0.25,  # Confidence threshold\n            iou=0.45,   # NMS IoU threshold\n            max_det=300,\n            device=0 if torch.cuda.is_available() else 'cpu',\n            verbose=False,\n            agnostic_nms=True,\n            retina_masks=True,\n        )\n        \n        # Convert results to submission format\n        batch_rows = convert_to_submission_format(results, batch_images)\n        all_submission_rows.extend(batch_rows)\n        \n    except Exception as e:\n        print(f\"Error in batch {i//INFERENCE_BATCH_SIZE}: {e}\")\n        continue\n\nprint(f\"\\nGenerated {len(all_submission_rows)} predictions\")\n\n# Create submission DataFrame\nprint(\"\\n=== Creating Submission File ===\")\nsubmission_df = pd.DataFrame(all_submission_rows)\n\n# Ensure all test images have at least one prediction\nall_test_ids = [os.path.basename(img).replace('.jpg', '') for img in test_images]\nif len(submission_df) > 0:\n    predicted_ids = submission_df['patient_id'].unique()\n    missing_ids = set(all_test_ids) - set(predicted_ids)\nelse:\n    missing_ids = set(all_test_ids)\n\n# Add dummy predictions for images without detections\nif missing_ids:\n    print(f\"Adding dummy predictions for {len(missing_ids)} images without detections\")\n    dummy_rows = []\n    for img_id in missing_ids:\n        # Create a small dummy polygon\n        dummy_rows.append({\n            'patient_id': img_id,\n            'class_id': 0,  # Default to class 0 (Abrasion)\n            'confidence': 0.01,  # Very low confidence\n            'poly': '0.1 0.1 0.1 0.2 0.2 0.2 0.2 0.1'  # Small square polygon\n        })\n    \n    submission_df = pd.concat([submission_df, pd.DataFrame(dummy_rows)], ignore_index=True)\n\n# Sort by patient_id and then by confidence (descending)\nsubmission_df = submission_df.sort_values(['patient_id', 'confidence'], ascending=[True, False])\n\n# Ensure correct column order\nsubmission_df = submission_df[['patient_id', 'class_id', 'confidence', 'poly']]\n\n# Save the main submission file\nsubmission_df.to_csv('submission.csv', index=False)\nprint(\"Main submission file created: submission.csv\")\n\n# Verify submission format\nprint(\"\\n=== Verifying Submission Format ===\")\nprint(f\"Total predictions: {len(submission_df)}\")\nprint(f\"Unique images: {submission_df['patient_id'].nunique()}\")\nprint(f\"All test images included: {submission_df['patient_id'].nunique() == len(test_images)}\")\n\n# Display first few rows\nprint(\"\\nFirst 5 rows of submission:\")\nprint(submission_df.head())\n\n# Check for any potential issues\nprint(\"\\n=== Checking for Potential Issues ===\")\n\n# Check for missing test images\nmissing_in_submission = set(all_test_ids) - set(submission_df['patient_id'].unique())\nif missing_in_submission:\n    print(f\"WARNING: Missing images in submission: {missing_in_submission}\")\nelse:\n    print(\"✓ All test images have predictions\")\n\n# Check class distribution\nprint(\"\\nPredictions per class:\")\nclass_dist = submission_df['class_id'].value_counts().sort_index()\nfor class_id, count in class_dist.items():\n    if 0 <= class_id < 9:\n        print(f\"  Class {class_id} ({CLASS_INFO[class_id]['name']}): {count}\")\n\n# Check confidence distribution\nprint(f\"\\nConfidence statistics:\")\nprint(f\"  Min: {submission_df['confidence'].min():.4f}\")\nprint(f\"  Max: {submission_df['confidence'].max():.4f}\")\nprint(f\"  Mean: {submission_df['confidence'].mean():.4f}\")\nprint(f\"  Median: {submission_df['confidence'].median():.4f}\")\n\n# Create alternative submission with higher confidence threshold\nprint(\"\\n=== Creating Alternative Submission (Higher Confidence) ===\")\nhigh_conf_df = submission_df[submission_df['confidence'] >= 0.3].copy()\n\n# Ensure all images still have at least one prediction\nhigh_conf_ids = high_conf_df['patient_id'].unique()\nmissing_high_conf = set(all_test_ids) - set(high_conf_ids)\n\nif missing_high_conf:\n    # Add the highest confidence prediction for each missing image\n    for img_id in missing_high_conf:\n        img_preds = submission_df[submission_df['patient_id'] == img_id]\n        if len(img_preds) > 0:\n            # Add the highest confidence prediction\n            high_conf_df = pd.concat([high_conf_df, img_preds.head(1)], ignore_index=True)\n        else:\n            # Add dummy prediction\n            dummy_row = pd.DataFrame([{\n                'patient_id': img_id,\n                'class_id': 0,\n                'confidence': 0.01,\n                'poly': '0.1 0.1 0.1 0.2 0.2 0.2 0.2 0.1'\n            }])\n            high_conf_df = pd.concat([high_conf_df, dummy_row], ignore_index=True)\n\nhigh_conf_df = high_conf_df.sort_values(['patient_id', 'confidence'], ascending=[True, False])\nhigh_conf_df.insert(0, 'id', range(1, len(high_conf_df) + 1))\nhigh_conf_df.to_csv('submission_high_conf.csv', index=False)\nprint(f\"Created high confidence submission with {len(high_conf_df)} predictions\")\n\nprint(\"\\n=== Pipeline Completed Successfully! ===\")\nprint(\"Submission files created:\")\nprint(\"  - submission.csv (main submission)\")\nprint(\"  - submission_high_conf.csv (alternative with higher confidence threshold)\")\nprint(\"\\nReady to submit to Kaggle!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T16:26:10.497153Z","iopub.execute_input":"2025-09-07T16:26:10.497424Z","iopub.status.idle":"2025-09-07T16:30:14.350939Z","shell.execute_reply.started":"2025-09-07T16:26:10.497402Z","shell.execute_reply":"2025-09-07T16:30:14.350179Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import cv2\nimport matplotlib.pyplot as plt\nfrom random import sample\nimport torch\n\n# Function to visualize predictions on images\ndef visualize_predictions(model, image_paths, class_names, num_samples=5, imgsz=640):\n    # Randomly pick some images\n    sample_paths = sample(image_paths, min(num_samples, len(image_paths)))\n\n    for img_path in sample_paths:\n        # Run prediction on the single image\n        results = model.predict(img_path, imgsz=imgsz, conf=0.25, iou=0.45, retina_masks=True, verbose=False)\n        result = results[0]\n\n        # Read image\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        # Draw boxes\n        if result.boxes is not None and len(result.boxes) > 0:\n            boxes = result.boxes.xyxy.cpu().numpy()  # x1, y1, x2, y2\n            classes = result.boxes.cls.cpu().numpy().astype(int)\n            confidences = result.boxes.conf.cpu().numpy()\n\n            for i, (box, cls, conf) in enumerate(zip(boxes, classes, confidences)):\n                x1, y1, x2, y2 = box.astype(int)\n                color = (0, 255, 0)  # green box\n                cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)\n                label = f\"{class_names[cls]} {conf:.2f}\"\n                # Increase font size\n                cv2.putText(\n                    img,\n                    label,\n                    (x1, y1-5),\n                    cv2.FONT_HERSHEY_SIMPLEX,\n                    fontScale=5.0,  # increase from 0.5\n                    color=color,\n                    thickness=10\n                )\n\n\n        if result.masks is not None and len(result.masks) > 0:\n            masks = result.masks.data.cpu().numpy()  # HxW mask\n            for mask in masks:\n                # Convert mask to binary and apply color overlay\n                mask = (mask > 0.5).astype(np.uint8)\n                color_mask = np.zeros_like(img)\n                color_mask[:, :, 0] = mask * 255  # Red channel\n                img = cv2.addWeighted(img, 1.0, color_mask, 0.5, 0)\n\n        plt.figure(figsize=(8, 8))\n        plt.imshow(img)\n        plt.axis('off')\n        plt.title(f\"Prediction for {os.path.basename(img_path)}\")\n        plt.show()\n\n\nvisualize_predictions(model, test_images, [CLASS_INFO[i]['name'] for i in range(4)], num_samples=5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T16:30:14.352254Z","iopub.execute_input":"2025-09-07T16:30:14.352498Z","iopub.status.idle":"2025-09-07T16:30:41.281719Z","shell.execute_reply.started":"2025-09-07T16:30:14.352479Z","shell.execute_reply":"2025-09-07T16:30:41.280896Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"yolo_model = model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T16:34:28.992246Z","iopub.execute_input":"2025-09-07T16:34:28.992575Z","iopub.status.idle":"2025-09-07T16:34:28.996703Z","shell.execute_reply.started":"2025-09-07T16:34:28.992553Z","shell.execute_reply":"2025-09-07T16:34:28.995872Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"OUTPUT_DIR = '/kaggle/working/'\n\nNEW_DATASET_PATH = f'{OUTPUT_DIR}/alphadent_90_10_split'\nNEW_TRAIN_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/train'\nNEW_EVAL_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/eval'\nNEW_VALID_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/valid'\nNEW_TEST_IMAGES_PATH = f'{NEW_DATASET_PATH}/images/test'\nNEW_TRAIN_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/train'\nNEW_EVAL_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/eval'\nNEW_VALID_LABELS_PATH = f'{NEW_DATASET_PATH}/labels/valid'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T16:34:31.033791Z","iopub.execute_input":"2025-09-07T16:34:31.034530Z","iopub.status.idle":"2025-09-07T16:34:31.038457Z","shell.execute_reply.started":"2025-09-07T16:34:31.034506Z","shell.execute_reply":"2025-09-07T16:34:31.037752Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Caries Classifier","metadata":{}},{"cell_type":"code","source":"import os\nimport cv2\nimport torch\nimport numpy as np\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import models, transforms\nfrom PIL import Image\nfrom glob import glob\n\n# ---------------- Paths & Settings ----------------\nTRAIN_PATH = \"/kaggle/input/alpha-dent/AlphaDent/images/train\"\nTRAIN_LABEL_PATH = \"/kaggle/input/alpha-dent/AlphaDent/labels/train\"\nNUM_CLASSES = 6   # caries1 -> caries6\nIMG_SIZE = 224    # ResNet input size\nBATCH_SIZE = 16\nEPOCHS = 5\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# ---------------- Dataset ----------------\nclass PolygonCariesDataset(Dataset):\n    def __init__(self, image_dir, label_dir, transform=None):\n        self.samples = []\n        self.transform = transform\n\n        image_paths = sorted(glob(os.path.join(image_dir, \"*.jpg\")))\n        label_paths = sorted(glob(os.path.join(label_dir, \"*.txt\")))\n\n        for img_path, lbl_path in zip(image_paths, label_paths):\n            with open(lbl_path, \"r\") as f:\n                lines = f.readlines()\n            for line in lines:\n                parts = line.strip().split()\n                class_id = int(parts[0])\n                if 3 <= class_id <= 8:  # merged classes\n                    coords = np.array(parts[1:], dtype=float).reshape(-1, 2)\n                    self.samples.append((img_path, coords, class_id-3))  # 0-5\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        img_path, coords, label = self.samples[idx]\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        h, w, _ = img.shape\n\n        # Crop polygon bounding box\n        xs, ys = coords[:,0], coords[:,1]\n        x1, y1 = max(int(xs.min()*w),0), max(int(ys.min()*h),0)\n        x2, y2 = min(int(xs.max()*w), w-1), min(int(ys.max()*h), h-1)\n        crop = img[y1:y2, x1:x2]\n\n        # Convert to PIL and transform\n        crop = Image.fromarray(crop).convert(\"RGB\")\n        if self.transform:\n            crop = self.transform(crop)\n        return crop, label\n\n# ---------------- Transforms ----------------\ntransform = transforms.Compose([\n    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.5]*3, [0.5]*3)\n])\n\nprint(\"--- Loading Train Dataset ---\")\ntrain_dataset = PolygonCariesDataset(TRAIN_PATH, TRAIN_LABEL_PATH, transform=transform)\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\n# ---------------- Model ----------------\nmodel = models.resnet18(pretrained=True)\nmodel.fc = nn.Linear(model.fc.in_features, NUM_CLASSES)  # Replace final layer\nmodel = model.to(DEVICE)\n\n# ---------------- Loss & Optimizer ----------------\ncriterion = nn.CrossEntropyLoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n\nprint(\"--- Start Training ---\")\n# ---------------- Training Loop ----------------\nfor epoch in range(EPOCHS):\n    model.train()\n    running_loss = 0.0\n    correct = 0\n    total = 0\n\n    for imgs, labels in train_loader:\n        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n\n        optimizer.zero_grad()\n        outputs = model(imgs)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n\n        running_loss += loss.item()\n        _, predicted = outputs.max(1)\n        total += labels.size(0)\n        correct += (predicted == labels).sum().item()\n\n    print(f\"Epoch [{epoch+1}/{EPOCHS}] Loss: {running_loss/len(train_loader):.4f} \"\n          f\"Accuracy: {correct/total:.4f}\")\n\n# ---------------- Save Model ----------------\ntorch.save(model.state_dict(), \"caries_polygon_classifier.pth\")\nprint(\"Model saved as caries_polygon_classifier.pth\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T16:35:20.018542Z","iopub.execute_input":"2025-09-07T16:35:20.019281Z","iopub.status.idle":"2025-09-07T17:24:43.677377Z","shell.execute_reply.started":"2025-09-07T16:35:20.019256Z","shell.execute_reply":"2025-09-07T17:24:43.676653Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Classifier Validation","metadata":{}},{"cell_type":"code","source":"from sklearn.metrics import classification_report, confusion_matrix\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport numpy as np\n\n# ---------------- Validation Dataset ----------------\nVAL_PATH = \"/kaggle/input/alpha-dent/AlphaDent/images/valid\"\nVAL_LABEL_PATH = \"/kaggle/input/alpha-dent/AlphaDent/labels/valid\"\n\nval_dataset = PolygonCariesDataset(VAL_PATH, VAL_LABEL_PATH, transform=transform)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n# ---------------- Evaluation ----------------\nmodel.eval()\nall_labels = []\nall_preds = []\n\nwith torch.no_grad():\n    for imgs, labels in val_loader:\n        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n        outputs = model(imgs)\n        _, preds = outputs.max(1)\n\n        all_labels.extend(labels.cpu().numpy())\n        all_preds.extend(preds.cpu().numpy())\n\n# ---------------- Metrics ----------------\nprint(\"\\n=== Classification Report ===\")\nprint(classification_report(all_labels, all_preds, target_names=[f\"caries{i}\" for i in range(1, NUM_CLASSES+1)]))\n\n# ---------------- Confusion Matrix ----------------\ncm = confusion_matrix(all_labels, all_preds)\nplt.figure(figsize=(8,6))\nsns.heatmap(cm, annot=True, fmt=\"d\", cmap=\"Blues\",\n            xticklabels=[f\"caries{i}\" for i in range(1, NUM_CLASSES+1)],\n            yticklabels=[f\"caries{i}\" for i in range(1, NUM_CLASSES+1)])\nplt.xlabel(\"Predicted\")\nplt.ylabel(\"True\")\nplt.title(\"Confusion Matrix\")\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T17:51:09.297620Z","iopub.execute_input":"2025-09-07T17:51:09.298107Z","iopub.status.idle":"2025-09-07T17:51:57.092516Z","shell.execute_reply.started":"2025-09-07T17:51:09.298084Z","shell.execute_reply":"2025-09-07T17:51:57.091931Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Classify Test Dataset","metadata":{}},{"cell_type":"code","source":"CLASS_INFO = {\n    0: {'name': 'Abrasion', 'description': 'Teeth with mechanical wear of hard tissues'},\n    1: {'name': 'Filling', 'description': 'Dental fillings of various types'},\n    2: {'name': 'Crown', 'description': 'Dental crown (restoration)'},\n    3: {'name': 'Caries 1', 'description': 'Caries in fissures and pits (occlusal surfaces of molars/premolars, buccal surfaces of molars, lingual surfaces of upper incisors).'},\n    4: {'name': 'Caries 2', 'description': 'Caries on proximal surfaces of molars and premolars.'},\n    5: {'name': 'Caries 3', 'description': 'Caries on proximal surfaces of incisors/canines without incisal edge involvement.'},\n    6: {'name': 'Caries 4', 'description': 'Caries on proximal surfaces of incisors/canines with incisal edge involvement.'},\n    7: {'name': 'Caries 5', 'description': 'Cervical caries (buccal/lingual surfaces).'},\n    8: {'name': 'Caries 6', 'description': 'Caries on incisal edges of anterior teeth or cusps of molars.'},\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T17:52:55.377889Z","iopub.execute_input":"2025-09-07T17:52:55.378681Z","iopub.status.idle":"2025-09-07T17:52:55.383242Z","shell.execute_reply.started":"2025-09-07T17:52:55.378655Z","shell.execute_reply":"2025-09-07T17:52:55.382530Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import torch\n\n# print(torch.cuda.device_count())  # should show 2\n# print(torch.cuda.get_device_name(0))  # GPU 0\n# print(torch.cuda.get_device_name(1))  # GPU 1\n\ntorch.cuda.empty_cache()\nbest_model_path = f'{OUTPUT_DIR}/alphadent_yolov8x_90_10/weights/best.pt'\nif os.path.exists(best_model_path):\n    yolo_model = YOLO(best_model_path)\n    print(f\"Loaded best model from: {best_model_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T18:15:44.583502Z","iopub.execute_input":"2025-09-07T18:15:44.584269Z","iopub.status.idle":"2025-09-07T18:15:44.707413Z","shell.execute_reply.started":"2025-09-07T18:15:44.584244Z","shell.execute_reply":"2025-09-07T18:15:44.706743Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import cv2\nimport torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport os\nfrom torchvision import transforms\nfrom PIL import Image\n\nclassifier = model\nclassifier.load_state_dict(torch.load(\"caries_polygon_classifier.pth\", map_location=\"cpu\"))\nclassifier.eval()\n\ndevice = torch.device(\"cuda:1\" if torch.cuda.is_available() else \"cpu\")\n# device = torch.device(\"cpu\")\nyolo_model = yolo_model.to(device)\nclassifier = classifier.to(device)\n\n# print(\"Using device:\", device, \"Name:\", torch.cuda.get_device_name(0))\n\n# ----- Transform for classifier -----\nclf_transform = transforms.Compose([\n    transforms.Resize((224, 224)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.5]*3, [0.5]*3)\n])\ndef predict_and_save(yolo_model, image_paths, class_names, submission_file=\"final_model_submission.csv\", imgsz=640):\n    all_preds = []\n    row_id = 0\n\n    # YOLO on GPU, classifier on CPU\n    classifier_cpu = classifier.to(\"cpu\")\n\n    for img_path in image_paths:\n        # YOLO runs normally (uses CUDA)\n        results = yolo_model.predict(\n            img_path, imgsz=imgsz, conf=0.25, iou=0.45,\n            retina_masks=True, verbose=False\n        )\n        result = results[0]\n\n        h, w = result.orig_shape\n        if result.boxes is None or len(result.boxes) == 0:\n            continue\n\n        boxes = result.boxes.xyxy.cpu().numpy()\n        classes = result.boxes.cls.cpu().numpy().astype(int)\n        confidences = result.boxes.conf.cpu().numpy()\n        masks_xy = result.masks.xy if result.masks is not None else []\n\n        for i, (box, cls, conf) in enumerate(zip(boxes, classes, confidences)):\n            x1, y1, x2, y2 = box.astype(int)\n\n            final_cls = int(cls)\n\n            # Refine \"caries merged\" class with CPU classifier\n            if cls == 3:\n                img_bgr = cv2.imread(img_path)\n                img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)\n                crop = img_rgb[y1:y2, x1:x2]\n\n                if crop.size > 0:\n                    crop_pil = Image.fromarray(crop).convert(\"RGB\")\n                    crop_tensor = clf_transform(crop_pil).unsqueeze(0)  # keep on CPU\n\n                    with torch.no_grad():\n                        out = classifier_cpu(crop_tensor)\n                        fine_cls = int(out.argmax(1).item())\n                    final_cls = fine_cls + 3  # shift to [3–8]\n\n            # polygon normalization\n            poly_str = \"\"\n            if i < len(masks_xy):\n                polygon = masks_xy[i]\n                if polygon is not None and len(polygon) >= 3:\n                    normalized_coords = []\n                    for point in polygon:\n                        x_norm = float(point[0]) / w\n                        y_norm = float(point[1]) / h\n                        normalized_coords.extend([\n                            max(0.0, min(1.0, x_norm)),\n                            max(0.0, min(1.0, y_norm))\n                        ])\n                    poly_str = \" \".join([f\"{coord:.6f}\" for coord in normalized_coords])\n\n            patient_id = os.path.splitext(os.path.basename(img_path))[0]\n            all_preds.append({\n                \"id\": row_id,\n                \"patient_id\": patient_id,\n                \"class_id\": final_cls,\n                \"confidence\": float(conf),\n                \"poly\": poly_str\n            })\n            row_id += 1\n\n    df = pd.DataFrame(all_preds)\n    df.to_csv(submission_file, index=False)\n    print(f\"✅ Submission saved to {submission_file}\")\n\n\npredict_and_save(\n    yolo_model=yolo_model,                        # your YOLO model\n    image_paths=test_images,                 # list of test images\n    class_names=[CLASS_INFO[i]['name'] for i in range(9)],  # 0–8 names\n    submission_file=\"submission.csv\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T18:16:26.130224Z","iopub.execute_input":"2025-09-07T18:16:26.130960Z","iopub.status.idle":"2025-09-07T18:21:15.517283Z","shell.execute_reply.started":"2025-09-07T18:16:26.130937Z","shell.execute_reply":"2025-09-07T18:21:15.516562Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport cv2\nimport torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom random import sample\nfrom PIL import Image\n\ndef visualize_predictions_with_classifier(\n    yolo_model,\n    image_paths,\n    class_names,\n    classifier,\n    clf_transform,\n    num_samples=5,\n    imgsz=640\n):\n    # Classifier on CPU\n    classifier = classifier.to(\"cpu\")\n    classifier.eval()\n\n    # Pick sample images\n    sample_paths = sample(image_paths, min(num_samples, len(image_paths)))\n\n    for img_path in sample_paths:\n        # Run YOLO prediction\n        results = yolo_model.predict(\n            img_path,\n            imgsz=imgsz,\n            conf=0.25,\n            iou=0.45,\n            retina_masks=True,\n            verbose=False\n        )\n        result = results[0]\n\n        # Read image\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        if result.boxes is not None and len(result.boxes) > 0:\n            boxes = result.boxes.xyxy.cpu().numpy()\n            classes = result.boxes.cls.cpu().numpy().astype(int)\n            confidences = result.boxes.conf.cpu().numpy()\n\n            for i, (box, cls, conf) in enumerate(zip(boxes, classes, confidences)):\n                x1, y1, x2, y2 = box.astype(int)\n\n                # Extract mask (if available)\n                mask_overlay = None\n                if result.masks is not None and len(result.masks.data) > i:\n                    mask = result.masks.data[i].cpu().numpy()\n                    mask = (mask > 0.5).astype(np.uint8) * 255\n                    mask_overlay = cv2.resize(mask, (img.shape[1], img.shape[0]))\n\n                # Fine classification if class == 3\n                if cls == 3:\n                    crop = img[y1:y2, x1:x2]\n                    if crop.size > 0:\n                        crop_pil = Image.fromarray(crop).convert(\"RGB\")\n                        crop_tensor = clf_transform(crop_pil).unsqueeze(0)  # CPU only\n                        with torch.inference_mode():\n                            pred = classifier(crop_tensor)\n                            fine_cls = pred.argmax(1).item()\n                        cls = fine_cls + 3\n\n                # Overlay mask with transparency\n                if mask_overlay is not None:\n                    color_mask = np.zeros_like(img, dtype=np.uint8)\n                    color_mask[:, :] = (0, 0, 200)  # green mask\n                    img = cv2.addWeighted(\n                        img, 1.0, cv2.bitwise_and(color_mask, color_mask, mask=mask_overlay), 0.4, 0\n                    )\n                \n                # Draw bounding box (thicker)\n                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 8)\n                \n                # Label (larger font + thicker text)\n                label = f\"{class_names[cls]} {conf:.2f}\"\n                cv2.putText(\n                    img, label, (x1, max(30, y1 - 10)),   # shift a bit higher\n                    cv2.FONT_HERSHEY_SIMPLEX,\n                    fontScale=4,    # larger text\n                    color=(0, 255, 0),\n                    thickness=10,      # thicker border\n                    lineType=cv2.LINE_AA\n                )\n\n\n        # Show final image\n        plt.figure(figsize=(8, 8))\n        plt.imshow(img)\n        plt.axis('off')\n        plt.title(f\"Prediction for {os.path.basename(img_path)}\")\n        plt.show()\n\n\nvisualize_predictions_with_classifier(\n    yolo_model=yolo_model,\n    image_paths=test_images,\n    class_names=[CLASS_INFO[i]['name'] for i in range(9)],  # 0–8 names\n    classifier=classifier,   # your trained classifier\n    clf_transform=clf_transform,\n    num_samples=5\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-07T19:07:18.295290Z","iopub.execute_input":"2025-09-07T19:07:18.295583Z","iopub.status.idle":"2025-09-07T19:07:45.392603Z","shell.execute_reply.started":"2025-09-07T19:07:18.295563Z","shell.execute_reply":"2025-09-07T19:07:45.392009Z"}},"outputs":[],"execution_count":null}]}