{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nfrom matplotlib import pyplot as plt\nimport cv2\nfrom PIL import Image\nimport torch\nimport torch.utils.data\nfrom torchvision.transforms import functional as F\nimport torchvision.transforms as transforms\nimport re\nfrom torch.utils.data import DataLoader, Dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\n#  将bbox的四个数转成x、y、w、h\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\n#  将bbox的四个数转成x、y、w、h\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\n# 丢弃bbox这一栏，换成用x、y、w、h\ntrain_df.drop(columns=['bbox'], inplace=True)\n# 边框转成float形式\ntrain_df['x'] = train_df['x'].astype(np.float)\ntrain_df['y'] = train_df['y'].astype(np.float)\ntrain_df['w'] = train_df['w'].astype(np.float)\ntrain_df['h'] = train_df['h'].astype(np.float)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = train_df['image_id'].unique() #去重\n# 划分数据集，先用id来划分\nvalid_ids = image_ids[-665:]\ntrain_ids = image_ids[:-665]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# 再根据id来划分\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# 继承了torch的Dataset\nclass WheatDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n        self.image_ids = dataframe['image_id'].unique() # 图像ID\n        self.df = dataframe #数据集\n        self.image_dir = image_dir\n        self.transforms = transforms #图像转化工具\n\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        # 读图片，并做一些转化\n        image = Image.open(f'{self.image_dir}/{image_id}.jpg').convert(\"RGB\")\n        image = np.array(image)\n        image = F.to_tensor(image)\n\n        # 处理一下边框，转换成左下和右上两个点\n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        \n        # 边框的面积\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n\n        # 所有标签都为1\n        # there is only one class\n        labels = torch.ones((records.shape[0],), dtype=torch.int64)\n        \n        # 这个是什么？\n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)\n        \n        # 将这张图片的信息放在target字典里面\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        return image, target\n\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# 这个函数要来干嘛？\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n# 训练集和验证集\ntrain_dataset = WheatDataset(train_df, DIR_TRAIN)\nvalid_dataset = WheatDataset(valid_df, DIR_TRAIN)\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=16,\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":{"trusted":true},"cell_type":"code","source":"import torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\n\n# our dataset has two classes only - background and person\nnum_classes = 2\n# get the model using our helper function\n\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\nin_features = model.roi_heads.box_predictor.cls_score.in_features\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n\n# move model to the right device\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\nmodel.to(device)\n\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# and a learning rate scheduler which decreases the learning rate by\n# 10x every 3 epochs\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_epochs = 20\n\nfor epoch in range(num_epochs): \n    sum_loss = 0.0\n    cal_freq = 10\n    itr = 1\n    for images, targets in train_data_loader: # 我猜就是一下子给batch_size个数据，迭代size/batch_size次\n        # 为什么要做这两步工作？ 为了.to(device)，转移到cuda:0上\n        # 获取图像本体和信息\n        images = list(image.to(device) for image in images)\n        # 是一个list，每个元素是一个dict，对应一张image\n        targets = [{k:v.to(device) for k,v in t.items()} for t in targets]\n\n        # 计算损失值\n        # 疑问1：损失函数是什么？在哪里设定？\n        # 疑问2：输入图像的大小？在哪里设定？\n        loss_dict = model(images, targets)\n\n        # 统计损失值之和，并梯度下降\n        losses = sum(loss for loss in loss_dict.values())\n        optimizer.zero_grad() # ？？？\n        losses.backward()  #反向传播\n        optimizer.step() # ？？？\n\n        # 用于输出训练过程 \n        sum_loss += losses.item() # 损失值累加\n        if itr % cal_freq == 0:\n            print(\"Iteration %d loss:%.03f\"%(itr,sum_loss/cal_freq))\n            sum_loss = 0\n        itr += 1\n\n    # update the learning rate\n    if lr_scheduler is not None:\n        lr_scheduler.step()\n    print('--------------Epoch %d finished'%epoch) ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"images, targets = next(iter(valid_data_loader))\nimages = list(image.to(device) for image in images)  #转到GPU上\ntargets = [{k: v.to(device) for k, v in t.items()} for t in targets]  #转到GPU上","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_index = 0\nboxes = targets[test_index]['boxes'].cpu().numpy().astype(np.int32)\n#  这个permute函数用来干什么？ 哦，原本通道是放到第一维的，现在放到第三维？\nin_sample = images[test_index]\nsample = in_sample.permute(1,2,0).cpu().numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# put the model in evaluation mode\nmodel.eval()\nwith torch.no_grad():\n    prediction = model([in_sample.to(device)])[0]\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n    for box in prediction['boxes']:\n        cv2.rectangle(sample,(box[0], box[1]),(box[2], box[3]),(220, 0, 0), 3)\nax.set_axis_off()\nax.imshow(sample)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}