{"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":[{"sourceId":19989,"databundleVersionId":1160143,"sourceType":"competition"}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torchvision\nimport numpy as np\nimport pandas as pd\nimport os\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n\nfrom torch.utils.data import Dataset, Subset, DataLoader\nfrom torchvision.ops import box_iou\nfrom torchmetrics.detection.mean_ap import MeanAveragePrecision\nimport torchvision.transforms as T\nfrom torchvision.datasets import ImageFolder\nfrom torch.utils.tensorboard import SummaryWriter\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nfrom torch import optim\n\nimport cv2","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:46:46.255424Z","iopub.execute_input":"2025-06-29T11:46:46.256073Z","iopub.status.idle":"2025-06-29T11:47:13.619131Z","shell.execute_reply.started":"2025-06-29T11:46:46.256041Z","shell.execute_reply":"2025-06-29T11:47:13.618379Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_csv = pd.read_csv(\"/kaggle/input/global-wheat-detection/train.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:13.620629Z","iopub.execute_input":"2025-06-29T11:47:13.621182Z","iopub.status.idle":"2025-06-29T11:47:13.862833Z","shell.execute_reply.started":"2025-06-29T11:47:13.621164Z","shell.execute_reply":"2025-06-29T11:47:13.862021Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_csv.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:13.863633Z","iopub.execute_input":"2025-06-29T11:47:13.863870Z","iopub.status.idle":"2025-06-29T11:47:13.884056Z","shell.execute_reply.started":"2025-06-29T11:47:13.863853Z","shell.execute_reply":"2025-06-29T11:47:13.883374Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_csv.info()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:13.884741Z","iopub.execute_input":"2025-06-29T11:47:13.885006Z","iopub.status.idle":"2025-06-29T11:47:13.924287Z","shell.execute_reply.started":"2025-06-29T11:47:13.884983Z","shell.execute_reply":"2025-06-29T11:47:13.923696Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_csv.groupby(\"image_id\")[[\"width\"]].apply(\"count\").hist(figsize=(10,5), bins=35)\nplt.title(\"Number of BBoxes\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:13.926167Z","iopub.execute_input":"2025-06-29T11:47:13.926388Z","iopub.status.idle":"2025-06-29T11:47:14.246025Z","shell.execute_reply.started":"2025-06-29T11:47:13.926371Z","shell.execute_reply":"2025-06-29T11:47:14.245383Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(train_csv.image_id.unique())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:14.246720Z","iopub.execute_input":"2025-06-29T11:47:14.246972Z","iopub.status.idle":"2025-06-29T11:47:14.258313Z","shell.execute_reply.started":"2025-06-29T11:47:14.246954Z","shell.execute_reply":"2025-06-29T11:47:14.257554Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import ast\ntype(ast.literal_eval(train_csv.iloc[0, 3]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:14.259044Z","iopub.execute_input":"2025-06-29T11:47:14.259287Z","iopub.status.idle":"2025-06-29T11:47:14.273331Z","shell.execute_reply.started":"2025-06-29T11:47:14.259263Z","shell.execute_reply":"2025-06-29T11:47:14.272756Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import ast, os\n# map image_id → list of bboxes\nbbox_dict = (\n    train_csv\n    .assign(bbox=lambda df: df.bbox.map(ast.literal_eval))\n    .groupby(\"image_id\")\n    .bbox\n    .apply(list)\n    .to_dict()\n)\n# list of all ids\nimage_ids = list(bbox_dict.keys())\ntrain_dir  = \"/kaggle/input/global-wheat-detection/train\"\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:14.274072Z","iopub.execute_input":"2025-06-29T11:47:14.274324Z","iopub.status.idle":"2025-06-29T11:47:16.238442Z","shell.execute_reply.started":"2025-06-29T11:47:14.274300Z","shell.execute_reply":"2025-06-29T11:47:16.237894Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def visualize(image_id):\n    if image_id not in bbox_dict:\n        print(f\"{image_id} not found\")\n        return\n    path = os.path.join(train_dir, f\"{image_id}.jpg\")\n    img  = cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)\n    for bbox in bbox_dict[image_id]:\n        xmin, ymin, w, h = map(int, bbox)\n        cv2.rectangle(img, (xmin,ymin), (xmin+w, ymin+h), (255,0,0), 2)\n    plt.imshow(img); plt.axis(\"off\")\n\nvisualize(\"00764ad5d\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:16.239244Z","iopub.execute_input":"2025-06-29T11:47:16.239524Z","iopub.status.idle":"2025-06-29T11:47:16.584423Z","shell.execute_reply.started":"2025-06-29T11:47:16.239502Z","shell.execute_reply":"2025-06-29T11:47:16.583674Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Create preprocessing Pipeline","metadata":{}},{"cell_type":"markdown","source":"    Here we will create a custom dataset to handle all preprocessing, such as converting coordinates","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, imgs_path, annot_dict, transform=None):\n        self.imgs_path = imgs_path\n        self.annot_dict = annot_dict\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.annot_dict.keys())\n\n    def __getitem__(self, idx):\n\n        image_id = list(self.annot_dict.keys())[idx]\n        boxes = np.array(self.annot_dict[image_id], dtype=np.float32)\n\n        # 2. convert to x1,y1,x2,y2\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]  # x2 = xmin + width\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3] # y2 = ymin + height\n        \n        image_path = os.path.join(self.imgs_path, f\"{image_id}.jpg\")\n        img = Image.open(image_path).convert('RGB')\n        if self.transform:\n            img = self.transform(img)\n        target = {\n            \"labels\": torch.ones((boxes.shape[0]), dtype=torch.int64),\n            \"boxes\": torch.tensor(boxes, dtype=torch.float64)\n        }\n        return img, target\n\ndef detection_collate(batch):\n    images, targets = zip(*batch)\n    return list(images), list(targets)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:49:24.943438Z","iopub.execute_input":"2025-06-29T11:49:24.943985Z","iopub.status.idle":"2025-06-29T11:49:24.950311Z","shell.execute_reply.started":"2025-06-29T11:49:24.943961Z","shell.execute_reply":"2025-06-29T11:49:24.949558Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ntrain_dataset = CustomDataset(train_dir, bbox_dict, transform=None)\ntrain_dataloader = DataLoader(train_dataset,  \n                              batch_size=1, \n                              shuffle=True,\n                             collate_fn = detection_collate,\n                             pin_memory = True if torch.cuda.is_available() else False)\n\n# Display image and label.\ntrain_features, train_labels = next(iter(train_dataloader))\nimg = train_features[0]\nlabel = train_labels[0]\nprint(label[\"labels\"].shape)\nprint(label[\"boxes\"].shape)\nplt.xticks([])\nplt.yticks([])\nplt.imshow(img)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:49:26.101561Z","iopub.execute_input":"2025-06-29T11:49:26.102294Z","iopub.status.idle":"2025-06-29T11:49:26.328709Z","shell.execute_reply.started":"2025-06-29T11:49:26.102268Z","shell.execute_reply":"2025-06-29T11:49:26.328010Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(train_dataloader)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:49:27.655313Z","iopub.execute_input":"2025-06-29T11:49:27.656067Z","iopub.status.idle":"2025-06-29T11:49:27.660445Z","shell.execute_reply.started":"2025-06-29T11:49:27.656045Z","shell.execute_reply":"2025-06-29T11:49:27.659869Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Faster RCNN arch with Resnet50 backbone","metadata":{}},{"cell_type":"code","source":"def evaluate_metrics(model, dataloader, device):\n    \"\"\"\n    Runs model on val set, returns:\n      - mean IoU (per-image, best-of-preds)\n      - mAP@0.5\n    \"\"\"\n    model.eval()\n    model.to(device)\n    \n    # meter for COCO‐style AP at IoU=0.5\n    metric = MeanAveragePrecision(\n        iou_type=('bbox',),\n        class_metrics=False\n    )\n    \n    total_iou = 0.0\n    img_count = 0\n    \n    with torch.no_grad():\n        for images, targets in dataloader:\n            # move targets to device; keep images as PIL or [C,H,W] tensors\n            imgs = [img.to(device) for img in images]\n            targs = [{k: v.to(device) for k,v in t.items()} for t in targets]\n            \n            # forward\n            outputs = model(imgs, targs)\n            \n            # 1) accumulate IoU\n            for gt, pred in zip(targs, outputs):\n                gt_boxes   = gt['boxes']            # [N_gt,4]\n                pred_boxes = pred['boxes']          # [N_pred,4]\n                \n                if gt_boxes.numel()==0 or pred_boxes.numel()==0:\n                    continue\n                ious = box_iou(pred_boxes, gt_boxes)       # [N_pred, N_gt]\n                best = ious.max(dim=1).values              # best IoU per pred\n                total_iou += best.mean().item()\n                img_count += 1\n            \n            # 2) update mAP metric\n            # predictions must be list of dicts with same keys\n            preds = [\n                {\n                    'boxes': out['boxes'].cpu(),\n                    'scores': out['scores'].cpu(),\n                    'labels': out['labels'].cpu()\n                }\n                for out in outputs\n            ]\n            gts = [\n                {\n                    'boxes': targ['boxes'].cpu(),\n                    'labels': targ['labels'].cpu()\n                }\n                for targ in targs\n            ]\n            metric.update(preds, gts)\n    \n    mean_iou = total_iou / img_count if img_count else 0.0\n    stats = metric.compute()\n    map50_95 = stats['map']\n    ap50 = stats['map_50']\n    ap75 = stats['map_75']\n    \n    print(f\"Eval → Mean IoU: {mean_iou:.4f}\")\n    print(f\"mAP@[.50:.95]: {map50_95:.4f}\")\n    print(f\"AP@0.50:      {ap50:.4f}\")\n    print(f\"AP@0.75:      {ap75:.4f}\")\n    model.train()\n    return mean_iou, ap50, map50_95\n\ndef train_model(model, train_loader, \n                val_loader, num_epochs, optimizer, \n                log_every=20, device=\"cpu\", writer=None):\n    model.to(device).train()\n    \n    for epoch in range(num_epochs):\n        epoch_loss = 0.0\n        print(f\"----Epoch: {epoch+1}----\")\n        \n        for batch_idx, (imgs, targets) in enumerate(train_loader, 1):\n            imgs = [img.to(device) for img in imgs]\n            targets = [{k:v.to(device) for k,v in t.items()} for t in targets]\n\n            \n            # forward + loss\n            loss_dict = model(imgs, targets)\n            loss = sum(loss_dict.values())\n            \n            \n            # backward + step\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n\n            epoch_loss += loss.item()\n            if batch_idx % log_every == 0:\n                print(f\"Epoch {epoch+1} | Batch {batch_idx} | Loss {loss:.4f}\")\n        avg_loss = epoch_loss / batch_idx\n        print(f\"Epoch {epoch+1} done | Avg Loss {avg_loss:.4f}\")\n        if writer:\n            writer.add_scalar(\"Train/AvgLoss\", avg_loss, epoch)\n\n        # ─── Validation ───────────────────────────────────\n        mean_iou, ap50, map50_95 = evaluate_metrics(model, val_loader, device)\n        writer.add_scalar(\"Val/Mean_IoU\",        mean_iou,   epoch)\n        writer.add_scalar(\"Val/AP@0.50\",         ap50,       epoch)\n        writer.add_scalar(\"Val/mAP@[.50:.95]\",   map50_95,   epoch)\n        \n\n    model.to(\"cpu\")\n        \n            \n\n\n                ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:49:29.972975Z","iopub.execute_input":"2025-06-29T11:49:29.973251Z","iopub.status.idle":"2025-06-29T11:49:29.985335Z","shell.execute_reply.started":"2025-06-29T11:49:29.973231Z","shell.execute_reply":"2025-06-29T11:49:29.984681Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, FasterRCNN_ResNet50_FPN_V2_Weights\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom sklearn.model_selection import train_test_split\nweights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT\nmodel = fasterrcnn_resnet50_fpn_v2(weights=weights)\ntransform = weights.transforms()\nsimple_tf = T.Compose([\n    T.ToTensor(),   # PIL → FloatTensor [0,1]\n])\n\nnum_classes = 2 # wheat + background\nin_features = model.roi_heads.box_predictor.cls_score.in_features # original input size\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) # change head to match num_classes\n\ntrain_ind, val_ind = train_test_split(range(len(bbox_dict)), test_size=0.1, random_state=42) \n\nfull_set  = CustomDataset(train_dir, bbox_dict, transform=simple_tf)\ntrain_set = Subset(full_set, train_ind)\nval_set   = Subset(full_set, val_ind)\n\ntrain_dataloader = DataLoader(train_set,  \n                              batch_size=8, \n                              shuffle=True,\n                             collate_fn = detection_collate,\n                             pin_memory = True if torch.cuda.is_available() else False)\nval_dataloader = DataLoader(val_set,  \n                            batch_size=8, \n                            shuffle=True,\n                            collate_fn = detection_collate,\n                            pin_memory = True if torch.cuda.is_available() else False)\n\n# Training Start\ntorch.manual_seed(42)\ndevice = torch.device(\"cuda:0\") if torch.cuda.is_available() else \"cpu\"\nn_epochs = 10\noptimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n\nwriter = SummaryWriter(\"runs/wheat_detector\")\n\ntrain_model(model, train_dataloader, val_dataloader, n_epochs, optimizer, log_every=50, device=device, writer=writer)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:49:32.234714Z","iopub.execute_input":"2025-06-29T11:49:32.235026Z","iopub.status.idle":"2025-06-29T11:50:46.196923Z","shell.execute_reply.started":"2025-06-29T11:49:32.235004Z","shell.execute_reply":"2025-06-29T11:50:46.196033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport matplotlib.patches as patches\n\ndef visualize_batch(images, outputs, class_names=None, threshold=0.5):\n    \"\"\"\n    Visualize a batch of images and their detection outputs.\n    \n    Args:\n      images (List[Tensor]): List of image tensors [C, H, W] in [0,1].\n      outputs (List[Dict]): List of model outputs with keys 'boxes', 'labels', 'scores'.\n      class_names (Dict[int, str], optional): Mapping label indices to names.\n      threshold (float, optional): Confidence threshold to filter boxes.\n    \"\"\"\n    batch_size = len(images)\n    cols = min(4, batch_size)\n    rows = (batch_size + cols - 1) // cols\n    \n    fig, axs = plt.subplots(rows, cols, figsize=(4 * cols, 4 * rows))\n    axs = axs.flatten() if batch_size > 1 else [axs]\n    \n    for idx, (img, output) in enumerate(zip(images, outputs)):\n        ax = axs[idx]\n        # convert tensor to numpy image\n        img_np = img.cpu().permute(1, 2, 0).numpy()\n        ax.imshow(img_np)\n        ax.axis('off')\n        # draw boxes above threshold\n        boxes = output['boxes'].cpu().detach().numpy()\n        scores = output['scores'].cpu().detach().numpy()\n        labels = output.get('labels')\n        if labels is not None:\n            labels = labels.cpu().numpy()\n        \n        for box, score, label in zip(boxes, scores, labels if labels is not None else []):\n            if score < threshold:\n                continue\n            x1, y1, x2, y2 = box\n            width, height = x2 - x1, y2 - y1\n            rect = patches.Rectangle(\n                (x1, y1), width, height, linewidth=2,\n                edgecolor='r', facecolor='none'\n            )\n            ax.add_patch(rect)\n            # label text\n            class_text = class_names[label] if class_names and label in class_names else str(label)\n            ax.text(x1, y1 - 5, f\"{score:.2f}\", \n                    color='white', fontsize=8, backgroundcolor='r')\n    \n    # hide any extra subplots\n    for j in range(idx + 1, len(axs)):\n        axs[j].axis('off')\n    \n    plt.tight_layout()\n    plt.savefig(\"test_classification.png\")\n    plt.show()\n\n# Example usage:\nmodel.to(device).eval()\nvisuals = []\nimages = next(iter(val_dataloader))[0]\nwith torch.no_grad():\n    for img in images:\n        out = model([img.to(device)])[0]      # one image → less mem\n        visuals.append({k: v.cpu() for k,v in out.items()})\n\n# now visuals is on CPU, safe to plot\nvisualize_batch(images, visuals, class_names={1:'wheat'}, threshold=0.8)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T11:47:19.414990Z","iopub.status.idle":"2025-06-29T11:47:19.415347Z","shell.execute_reply.started":"2025-06-29T11:47:19.415229Z","shell.execute_reply":"2025-06-29T11:47:19.415241Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Adding Augmentations and Retrain","metadata":{}},{"cell_type":"markdown","source":"    The Albumentations library will be used, since it provides automatic transformations of bounding boxes","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, imgs_path, annot_dict, ids, transform):\n        self.imgs_path = imgs_path\n        self.annot_dict = {k:v for k,v in annot_dict.items() if k in ids}\n        self.transform = transform\n        self.ids = ids\n        \n    def __len__(self):\n        return len(self.annot_dict.keys())\n\n    def __getitem__(self, idx):\n        \n        image_id = self.ids[idx]\n        boxes = np.array(self.annot_dict[image_id], dtype=np.float32)\n\n        # 2. convert to x1,y1,x2,y2\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]  # x2 = xmin + width\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3] # y2 = ymin + height\n        \n        img = cv2.imread(f\"{self.imgs_path}/{image_id}.jpg\")\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        labels = [1] * len(boxes)\n\n        augmented = self.transform(image=img, bboxes=boxes, labels=labels)\n        img = augmented[\"image\"].float()\n        boxes = torch.tensor(augmented[\"bboxes\"], dtype=torch.float32)\n        labels = torch.tensor(augmented[\"labels\"], dtype=torch.int64)\n\n        target = {\"boxes\": boxes, \"labels\": labels}\n        return img, target\n\ndef detection_collate(batch):\n    images, targets = zip(*batch)\n    return list(images), list(targets)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T12:05:13.222534Z","iopub.execute_input":"2025-06-29T12:05:13.222847Z","iopub.status.idle":"2025-06-29T12:05:13.230117Z","shell.execute_reply.started":"2025-06-29T12:05:13.222825Z","shell.execute_reply":"2025-06-29T12:05:13.229354Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\n# Setting up architecture\nweights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT\nmodel_augmented = fasterrcnn_resnet50_fpn_v2(weights=weights)\ntrain_tf = A.Compose([\n    A.HorizontalFlip(p=0.5),\n    A.RandomResizedCrop((512,512)),\n    A.RandomBrightnessContrast(),\n    A.Resize(512, 512),\n    ToTensorV2(), # converts to [C,H,W] float tensor\n],\n    bbox_params=A.BboxParams(  \n      format='pascal_voc',  # [x_min, y_min, x_max, y_max]\n      label_fields=['labels']\n    ))\n\nval_tf = A.Compose([\n    A.Resize(512,512),\n    ToTensorV2(),\n],\n    bbox_params = A.BboxParams(\n        format=\"pascal_voc\",\n        label_fields=[\"labels\"]\n    ))\n\n\nnum_classes = 2 # wheat + background\nin_features = model_augmented.roi_heads.box_predictor.cls_score.in_features # original input size\nmodel_augmented.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) # change head to match num_classes\n\ntrain_ids, val_ids = train_test_split(list(bbox_dict.keys()), test_size=0.1, random_state=42) \n\ntrain_set = CustomDataset(train_dir, bbox_dict, train_ids, train_tf)\nval_set = CustomDataset(train_dir, bbox_dict, val_ids, val_tf)\n\n\ntrain_dataloader = DataLoader(train_set,  \n                              batch_size=8, \n                              shuffle=True,\n                             collate_fn = detection_collate,\n                              num_workers=4,\n                             pin_memory = True if torch.cuda.is_available() else False)\nval_dataloader = DataLoader(val_set,  \n                            batch_size=8, \n                            shuffle=True,\n                            collate_fn = detection_collate,\n                            num_workers=4,\n                            pin_memory = True if torch.cuda.is_available() else False)\n\n# Training Start\ntorch.manual_seed(42)\ndevice = torch.device(\"cuda:0\") if torch.cuda.is_available() else \"cpu\"\nn_epochs = 10\noptimizer = optim.SGD(model_augmented.parameters(), lr=0.001, momentum=0.9)\n\nwriter = SummaryWriter(\"runs/wheat_detector\")\n\ntrain_model(model_augmented, train_dataloader, val_dataloader, n_epochs, optimizer, log_every=50, device=device, writer=writer)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-29T12:09:06.577864Z","iopub.execute_input":"2025-06-29T12:09:06.578594Z","iopub.status.idle":"2025-06-29T13:26:24.998046Z","shell.execute_reply.started":"2025-06-29T12:09:06.578562Z","shell.execute_reply":"2025-06-29T13:26:24.997089Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}