{"cells":[{"metadata":{},"cell_type":"markdown","source":"# <center><h1>Wheat Spike Detection<h1></center>\n\n![](https://cdn.glutenfreeliving.com/2015/01/wheat-starch-image-825x338.jpg)\n\nMachine learning has a lot of applications in various industries. The recent development in technology has also enabled ML to step into the realm of agriculture. Image analysis has significantly enhanced the potential for achieving high-throughput analysis of crop fields. And has also enabled in detecting diseases in the crops at a very initial stage.  For wheat breeding purposes, assessing the production of wheat spikes, as the grain-bearing organ, is a useful proxy measure of grain production. Thus, being able to detect and characterize spikes from images of wheat fields is an essential component in a wheat breeding pipeline for the selection of high yielding varieties.","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 5GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# import required libraries\nimport numpy as np\n\n\nimport PIL\nfrom PIL import Image\n\n# plotly libraries\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nimport plotly.graph_objs as go\nfrom plotly.offline import iplot, init_notebook_mode\nimport plotly.express as px\nimport plotly.io as pio\n\nimport cv2","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# initialize source paths\nsource_path = '/kaggle/input/global-wheat-detection/'\ntrain_path = source_path + 'train/'\ntest_path = source_path + 'test/'\n\n# read data table\ntrain_df = pd.read_csv(source_path +'train.csv')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# <center><h1>Basic Data Exploration<h1></center>\n    \n    \n First we will look at the data source to get an intial understanding of what is there. It will help us to know how many images are there, what are the sources for the data, image sizes, bounding box information etc.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"As a first step we will create the individual bbox columns and also create the bounding box area column. It will be helpful if we take care of it now since we have to analyze this later.\n","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# expand the bbox column into seprate columns\ntrain_df[['bbox_xmin','bbox_ymin','bbox_width','bbox_height']] = train_df['bbox'].str.split(',',expand=True)\ntrain_df['bbox_xmin'] = train_df['bbox_xmin'].str.replace('[','').astype(float)\ntrain_df['bbox_ymin'] = train_df['bbox_ymin'].str.replace(' ','').astype(float)\ntrain_df['bbox_width'] = train_df['bbox_width'].str.replace(' ','').astype(float)\ntrain_df['bbox_height'] = train_df['bbox_height'].str.replace(']','').astype(float)\n\n# add xmax, ymax, and area columns for bounding box\ntrain_df['bbox_xmax'] = train_df['bbox_xmin'] + train_df['bbox_width']\ntrain_df['bbox_ymax'] = train_df['bbox_ymin'] + train_df['bbox_height']\ntrain_df['bbox_area'] = train_df['bbox_height'] * train_df['bbox_width']","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## How are the images distributed by source?","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# count distinct images by source\nimg_source_dist = train_df.groupby(['source']).agg(image_count=('image_id','nunique'),wheat_head=('image_id','size'))\nimg_source_dist.reset_index(inplace=True,drop=False)\nfig = px.pie(img_source_dist, values='image_count', names='source', title='Spike Distribution by Source')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## How about 'em wheat heads?","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"img_source_dist['Avg_Wheat_Head'] = img_source_dist['wheat_head']/img_source_dist['image_count']\nimg_source_dist = img_source_dist.sort_values(by='Avg_Wheat_Head', ascending=True)\n\nfig = go.Figure(data=[\n    go.Bar(name='Avg Wheat Head Count', x=img_source_dist['Avg_Wheat_Head'], y=img_source_dist['source'],\n           orientation='h',marker_color='salmon')\n])\n# Change the bar mode\nfig.update_layout(title_text='Avg Number of Spikes',\n                  height=400)\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The ethz_1 source has a significant number of wheat heads per image, approximately 69 per image. Maybe these are very small and will be highly concentrated. We will know better once we look at the images. Now let's look at the overall distribution for the wheat heads by image.","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"wheat_heads_per_image = train_df.groupby('image_id').agg(head_count=('image_id','size'))\nwheat_heads_per_image.reset_index(inplace=True, drop=False)\n\nfig = px.histogram(wheat_heads_per_image, x=\"head_count\",marginal=\"box\")\nfig.update_layout(\n    xaxis = dict(\n        title_text = \"Spike Count\"), title = 'Spike Count per Image')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The IQR is between 28 and 59 wheat heads per image, with a median value of 43. There are some images with a counts greater than 100 and as less as 1. We will look at images in these 3 ranges.","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# create list of images per the regions identified \nheads_large =  list(wheat_heads_per_image[wheat_heads_per_image.head_count > 100]['image_id'].unique())\nheads_normal = list(wheat_heads_per_image[(wheat_heads_per_image.head_count >= 30) & (wheat_heads_per_image.head_count <= 30)]['image_id'].unique())\nheads_small =  list(wheat_heads_per_image[wheat_heads_per_image.head_count <= 5]['image_id'].unique())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# define a function to display the images\ndef get_bbox(df, image_id):\n    bboxes = []\n    image_bbox = df[df.image_id == image_id]\n    \n    for _,row in image_bbox.iterrows():\n        bboxes.append([row.bbox_xmin, row.bbox_ymin, row.bbox_width, row.bbox_height])\n        \n    return bboxes\n        \ndef plot_image(images, title=None):\n    fig = plt.figure(figsize = (20,10))\n    for i in range(1,4):\n        ax = fig.add_subplot(1, 3, i)\n        img = np.random.choice(images)\n        image_path = os.path.join(train_path,img +'.jpg')\n        image = Image.open(image_path)\n        ax.imshow(image)\n    \n        b = get_bbox(train_df,img)\n    \n        for bbox in b:\n                    rect = patches.Rectangle((bbox[0],bbox[1]),bbox[2],bbox[3],linewidth=2,edgecolor='yellow',facecolor='none')\n                    ax.add_patch(rect)\n    plt.suptitle(title, fontsize=14)\n    plt.tight_layout()","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"plot_image(heads_small,'Spike Count <= 5')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"For the images with less than 5 wheat heads we see these are very small heads. And Some of them are on the ground and a few are not clearly visbible. We also have heads with the shoe background. \n\nAlso an other thing to note is there is a clear demarcation in the brightness of the pictures. Some of them are dark and some of them are light. We will have to use some augmentation techniques for these.","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"plot_image(heads_normal,'Spike Count >= 30 & <= 60')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The images with count between 30 and 60 look good. The images seem sharper, but we need to look at more of these and also analyze those parameters.","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"plot_image(heads_large,'Spike Count > 100')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Wow!! those look cluttered , but they are a good source of information.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Looking into the box...\nIn this section we will look at the bounding boxes. Specifically, the area distribution and how many of these narrow boxes we have in the dataset. The area becomes important for IOU detection.","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"fig = px.histogram(train_df, x=\"bbox_area\",marginal=\"box\")\nfig.update_layout(\n    xaxis = dict(\n        title_text = \"Bounding Box Area\"), title = 'Bounding Box Area Distribution')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The bounding box area distribution is extremely skewed. We have some exteremly large bounding boxes with area greater than 100,000. The maximum area is approximately 500,000. As stated earlier, these will have an impact on the IOU and may have to be dealt with.\n\nThe smallest are of the bounding box is 2.0. This is could be one of those wheat heads that is still forming. Nevertheless let's a take a look at some of the images with both types of bounding boxes.","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"large_area = list(train_df[train_df.bbox_area > 100000]['image_id'].unique())\nsmall_area = list(train_df[train_df.bbox_area <= 10]['image_id'].unique())","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"plot_image(large_area, title='Large Bounding Boxes')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"There are some really large boxes that don't make a lot sense!","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_image(small_area, title='Small Bounding Boxes')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# <center><h1>Image Augmentation<h1></center>\n    \nWe have only about 3,300 images in the training data. Deep neural networks need a lot of data to be effective. That is where image augmentation comes into play. It is the process of creating more images from the existing training data by applying transformations. These include, but not limited to, flips, adding blur, increase sharpness and more. Some of these are very helpful to increase the accuracy of the models.\n\nThere is a fantastic library called 'Albumentations' that helps in creating augmenations quickly and effectively within a few lines of code. The github link is [here](https://github.com/albumentations-team/albumentations) and is a really good resource for beginners. The library was created by Kaggle Grandmasters and has helped win Kaggle competitions.\n\nSince this is my first Kaggle competition and I have not really done object detection analysis, I am going to use the Albumentation library to play around a little bit to understand the different features available.\n","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import albumentations\nfrom albumentations import (\n    HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,\n    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,\n    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, RandomBrightnessContrast, IAAPiecewiseAffine,\n    IAASharpen, IAAEmboss, Flip, OneOf, Compose,VerticalFlip,BboxParams,Rotate, ChannelShuffle, RandomRain)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# define functions for augmentation and display\n\ndef get_aug(aug, min_area=0., min_visibility=0.):\n    return Compose(aug, bbox_params=BboxParams(format='pascal_voc', min_area=min_area, \n                                               min_visibility=min_visibility, label_fields=['labels']))\n\nBOX_COLOR = (255,255,0)\ndef visualize_bbox(img, bbox, color=BOX_COLOR, thickness=2):\n    x_min, y_min, x_max, y_max = bbox\n    x_min, x_max, y_min, y_max = int(x_min), int(x_max), int(y_min), int(y_max)\n    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness)   \n    return img\n\ndef visualize(annotations):\n    img = annotations['image'].copy()\n    for idx, bbox in enumerate(annotations['bboxes']):\n        img = visualize_bbox(img, bbox)\n    return img \n    \ndef aug_plots(image_id, aug, title=None):\n    img_path = os.path.join(train_path,image_id +'.jpg')\n    image = cv2.imread(img_path, cv2.IMREAD_COLOR)\n    bbox = train_df[train_df['image_id'] == image_id][['bbox_xmin', 'bbox_ymin', 'bbox_xmax', 'bbox_ymax']].astype(np.int32).values\n\n    labels = np.ones((len(bbox), ))\n    annotations = {'image': image, 'bboxes': bbox, 'labels': labels}\n    \n    aug = get_aug(aug)\n    augmented = aug(**annotations)\n    visualize(augmented)\n    \n    fig = plt.figure(figsize = (15,7))\n    ax = fig.add_subplot(1, 2, 1)\n    ax.imshow(visualize(annotations))\n    plt.title('Original')\n    \n    ax = fig.add_subplot(1, 2, 2)\n    ax.imshow(visualize(augmented))\n    plt.title(title)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_plots('b6ab77fd7',[VerticalFlip(p=1)], 'Vertical Flip')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_plots('69fc3d3ff',[HorizontalFlip(p=1)], 'Horizontal Flip')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_plots('69fc3d3ff',[Blur(blur_limit= 7,p=0.5)], 'Blur')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_plots('69fc3d3ff',[Rotate(p=0.5)], 'Rotate')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_plots('69fc3d3ff',[HueSaturationValue(hue_shift_limit=20, sat_shift_limit=50, val_shift_limit=50, p=1)], 'Hue Saturation')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This is a good one. The spikes stand out!","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_plots('69fc3d3ff',[ChannelShuffle(p=1)], 'Channel Shuffle')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_plots('69fc3d3ff',[GaussNoise()], 'Gauss Noise')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_plots('69fc3d3ff',[RandomRain(p=1, brightness_coefficient=0.9, drop_width=1, blur_value=5)], 'Random Rain')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This is really cool...adding weather aspects to the data.\n\nI can keep playing with this all day :-) But it is time to move on!","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# <center><h1>Train Model<h1></center>\n    \nWe will train a FasterRCNN model for this analysis. I am very new to DL and object detection. I am going to try my best to explain the proces in the most layman terms.\n\nR-CNN extracts a bunch of regions from the given image using selective search, and then checks if any of these boxes contains an object. We first extract these regions, and for each region, CNN is used to extract specific features. Finally, these features are then used to detect objects. Unfortunately, R-CNN becomes rather slow due to these multiple steps involved in the process.\n\nFast R-CNN, on the other hand, passes the entire image to ConvNet which generates regions of interest (instead of passing the extracted regions from the image). Also, instead of using three different models (as we saw in R-CNN), it uses a single model which extracts features from the regions, classifies them into different classes, and returns the bounding boxes.\n\nAll these steps are done simultaneously, thus making it execute faster as compared to R-CNN. Fast R-CNN is, however, not fast enough when applied on a large dataset as it also uses selective search for extracting the regions.\n\nFaster R-CNN fixes the problem of selective search by replacing it with Region Proposal Network (RPN). We first extract feature maps from the input image using ConvNet and then pass those maps through a RPN which returns object proposals. Finally, these maps are classified and the bounding boxes are predicted.\n\nBelow steps sumarize the steps:\n\n* Take an input image and pass it to the ConvNet which returns feature maps for the image\n* Apply Region Proposal Network (RPN) on these feature maps and get object proposals\n* Apply ROI pooling layer to bring down all the proposals to the same size\n* Finally, pass these proposals to a fully connected layer in order to classify any predict the bounding boxes for the image\n\nThe Faster R-CNN model will be implemented with PyTorch. The first step is to write a function to prepare the dataset.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The data prep step in PyTorch requires three main parameters:\n* an image with size h and w. The minimum size required is a 800 x 800 image.\n* a target dictionary that needs the following mandatory fields:\n        1. coordinates of the bounding boxes\n        2. labels for each bounding box - background is always 0.\n        3. image id - unique identifier\n        4. area - area of the bounding box\n        5. iscrowd - instacnces with iscrowd = True will be ignored (i don't know what this means and have set it to False\n        \nIn addition we can add masks if available and also specify transformations. \n\nThe function below will do all of the above.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        boxes = records[['bbox_xmin', 'bbox_ymin', 'bbox_xmax', 'bbox_ymax']].values\n        \n        area = records['bbox_area'].values  # i already have the area in my dataframe\n        area = torch.as_tensor(area, dtype=torch.float32)\n\n        # there is only one class - so all will be 1\n        labels = torch.ones((records.shape[0],), dtype=torch.int64) \n        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['image_id'] = torch.tensor([index])\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n\n        if self.transforms:\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n            \n            target['boxes'] = torch.tensor(sample['bboxes'])\n            target['boxes'] = target['boxes'].type(torch.float32)\n\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We have already seen that the albumentations library is a good resource. So we will start only with a flip for the train. Note that the validation set should not be augmented!","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# define transformation functions\ndef get_train_transform():\n    return A.Compose([\n        A.Flip(0.5),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"There are two common situations where one might want to modify one of the available models in torchvision modelzoo. The first is when we want to start from a pre-trained model, and just finetune the last layer. The other is when we want to replace the backbone of the model with a different one (for faster predictions, for example).\n\nIn this case we will use the pre-trained model and finetune the last layer since our dataset is not that large.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_model_instance_segmentation(num_classes):\n    # load an instance segmentation model pre-trained pre-trained on COCO\n    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n\n    # get number of input features for the classifier\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    # replace the pre-trained head with a new one\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Next, we split the data into a train_set (2,708 images) and a validation set (665 images). ","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = train_df['image_id'].unique()\nvalid_ids = image_ids[-665:]\ntrain_ids = image_ids[:-665]\n\nval_set = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_set = train_df[train_df['image_id'].isin(train_ids)]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Next, we define the dataloaders. These steps are pretty self explanatory.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_dataset = WheatDataset(train_df, train_path, get_train_transform())\nvalid_dataset = WheatDataset(valid_df, train_path, get_valid_transform())\n\n\n# split the dataset in train and test set\nindices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now, we setup the model.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# train on the GPU or on the CPU, if a GPU is not available\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n# define the number of classes\nnum_classes = 2 # one for wheat and one for background\n\n# get the model using our helper function\nmodel = get_model_instance_segmentation(num_classes)\n\n# move model to the right device\nmodel.to(device)\n\n# construct an optimizer\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005,momentum=0.9, weight_decay=0.0005) \n\n\nnum_epochs = 4\n\n#for epoch in range(num_epochs):\n    # train for one epoch, printing every 10 iterations\n    #train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=10)\n    # update the learning rate\n    #lr_scheduler.step()\n    # evaluate on the test dataset\n   # evaluate(model, valid_data_loader, device=device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from engine import train_one_epoch","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}