{"cells":[{"metadata":{"_uuid":"596fb18d-7de6-4aa7-bde8-f01b37d3c932","_cell_guid":"83fe3f4b-3611-4531-bec7-faedc16bddbf","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\nfrom sklearn.utils import shuffle\nimport random\n# for visualize\nimport copy\n\n# from PIL import Image\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n# from albumentations.pytorch.transforms import ToTensorV2, ToTensor\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\n\nfrom matplotlib import pyplot as plt\ngodimg = 0\n\nDIR_INPUT = '../input'\nDIR_TRAIN = os.path.join(DIR_INPUT, 'global-wheat-detection/train')\nDIR_TEST = os.path.join(DIR_INPUT, 'global-wheat-detection/test')\n\nTRAIN_CSV = os.path.join(DIR_INPUT, 'global-wheat-detection/train.csv')\nSAMPLE_SUBMIT_CSV = os.path.join(DIR_INPUT, 'global-wheat-detection/sample_submission.csv')\n\ntrain_df = pd.read_csv(TRAIN_CSV)\ntest_df = pd.read_csv(SAMPLE_SUBMIT_CSV)\ntrain_df.shape\n\ntrain_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\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\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\ntrain_df.drop(columns=['bbox'], inplace=True)\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)\n\nimage_ids = train_df['image_id'].unique()\nvalid_ids = image_ids[-10:]\ntrain_ids = image_ids[:-10]\n\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]\n\n\nclass WheatDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, stage_flag, transforms=None):\n        super().__init__()\n        \n        self.df = dataframe\n        self.image_ids = dataframe['image_id'].unique()\n        self.image_ids = shuffle(self.image_ids)\n        self.labels = [np.zeros((0, 5), dtype=np.float32)] * len(self.image_ids)\n        self.img_size = 1024\n        im_w = 1024\n        im_h = 1024\n        for i, img_id in enumerate(self.image_ids):\n            records = self.df[self.df['image_id'] == img_id]\n            boxes = records[['x', 'y', 'w', 'h']].values\n            boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n            boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n            \n            label_tmp = []\n            for box in boxes:\n                x1, y1, x2, y2 = box\n                label_tmp.append([0, x1, y1, x2, y2])\n            self.labels[i] = np.array(label_tmp)\n        \n        self.image_dir = image_dir\n        self.transforms = transforms\n        \n        # self.mosaic = False\n        \n        self.augment = True\n        if stage_flag == \"val\":\n            self.augment = False\n    \n    def load_image(self, index):\n        # loads 1 image from dataset, returns img, original hw, resized hw\n        img = cv2.imread( os.path.join(DIR_TRAIN, str(self.image_ids[index])+\".jpg\"), cv2.IMREAD_COLOR)\n        assert img is not None, 'Image Not Found ' + DIR_TRAIN\n        h0, w0 = img.shape[:2]  # orig hw\n        return img, (h0, w0), img.shape[:2]  # img, hw_original, hw_resized\n    \n    def vis_show(self, img, index, labels):\n        # print(img.shape)\n        img_show = copy.deepcopy(img)\n        for box in labels.astype(int):\n            cv2.rectangle(img_show, (box[1], box[2]), (box[3], box[4]),(0,255,0), 2)\n        save_path = './tmp_show_aug'\n        if not os.path.exists(save_path):\n            os.makedirs(save_path)\n        cv2.imwrite(os.path.join(save_path, self.image_ids[index]+'.png'), img_show)\n    \n    def __getitem__(self, index: int):\n        # Load image\n        img, (h0, w0), (h, w) = self.load_image(index)\n        labels = self.labels[index]\n        # self.vis_show(img, index, labels)\n        \n        if self.augment:\n            aug_methods = [\"flip\", \"rotate\",\"mosaic\",\"light_contrast\"]\n            #print(\"---aug_methods num: \", len(aug_methods))\n            aug_num = 5\n            while aug_num > 0:\n                aug_ind = random.randint(0, len(aug_methods)-1) #[0, len(aug_methods)-1]\n                if aug_methods[aug_ind]==\"flip\":\n                    img, labels = augment_flip(img, labels)\n                elif aug_methods[aug_ind]==\"rotate\":\n                    if random.randint(0,1)==1:\n                        degree = 90\n                    else:\n                        degree = -90\n                    img, labels = random_rotate(img, labels, degrees=degree)\n                elif aug_methods[aug_ind]==\"affine\":\n                    # Augment imagespace\n                    degree = random.randint(-5, 5)\n                    shear = random.randint(-5, 5)\n                    img, labels = random_affine(img, labels, degrees=degree, translate=0, scale=0, shear=shear)\n                elif aug_methods[aug_ind]==\"hsv\":\n                    # Augment colorspace\n                    # augment_hsv(img, hgain=0.0138, sgain= 0.678, vgain=0.36)\n                    img = augment_hsv2(img)\n                elif aug_methods[aug_ind]==\"mixup\":\n                    # Augment mixup_image\n                    index_r = random.randint(0, self.image_ids.shape[0] - 1)\n                    img_r, _, _ = self.load_image(index_r)\n                    labels_r = self.labels[index_r]\n                    img, labels = augment_mixup(img, labels, img_r, labels_r, alpha=0.5)\n                    # self.vis_show(img, index, labels)\n                elif aug_methods[aug_ind]==\"mosaic\":\n                    index_1 = random.randint(0, self.image_ids.shape[0] - 1)\n                    img_1, _, _ = self.load_image(index_1)\n                    labels_1 = self.labels[index_1]\n                    index_2 = random.randint(0, self.image_ids.shape[0] - 1)\n                    img_2, _, _ = self.load_image(index_2)\n                    labels_2 = self.labels[index_2]\n                    index_3 = random.randint(0, self.image_ids.shape[0] - 1)\n                    img_3, _, _ = self.load_image(index_3)\n                    labels_3 = self.labels[index_3]\n                    img, labels = augment_mosaic(img, labels, img_1, labels_1, img_2, labels_2, img_3, labels_3)\n                elif aug_methods[aug_ind]==\"light_contrast\":\n                    img = augment_light_contrast(img)\n                elif aug_methods[aug_ind]==\"illumination\":\n                    scale = 70.0\n                    normallization = True\n                    img = augment_illumination(img, scale, normallization)\n                aug_num-=1\n        \n        # if labels[:,0].shape[0] == 0:\n            # print()\n        #self.vis_show(img, index, labels)\n        \n        d = {}\n        d['boxes'] = torch.from_numpy(labels[:,1:].astype(np.float32))\n        d['labels'] = torch.ones((labels[:,0].shape[0],), dtype=torch.int64)\n        \n        #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)\n        #img /= 255.0\n        #return torch.from_numpy(img), d\n        return torch.from_numpy(torch.from_numpy(img).permute(2, 0, 1).numpy().astype(np.float32) / 255.0), d\n \n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n    \n    # def load_image_and_boxes(self, index):\n        # image_id = self.image_ids[index]\n        # image = cv2.imread( os.path.join(DIR_TRAIN, str(image_id)+\".jpg\"), cv2.IMREAD_COLOR)\n        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        # image /= 255.0\n        # records = self.df[self.df['image_id'] == image_id]\n        # boxes = records[['x', 'y', 'w', 'h']].values\n        # boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        # boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        # return image, boxes\n     \nclass WheatTestDataset(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(os.path.join(self.image_dir, str(image_id)+\".jpg\"), cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        if self.transforms:\n            sample = {\n                'image': image,\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n\n        return image, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n        \nclass Averager:\n    def __init__(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n    def send(self, value):\n        self.current_total += value\n        self.iterations += 1\n\n    @property\n    def value(self):\n        if self.iterations == 0:\n            return 0\n        else:\n            return 1.0 * self.current_total / self.iterations\n\n    def reset(self):\n        self.current_total = 0.0\n        self.iterations = 0.0        \n\n#######################################################################################\ndef augment_mosaic(img1, label1, img2, label2, img3, label3, img4,label4):\n    h = img1.shape[0]\n    w = img1.shape[1]\n    cut_x = random.randint(int(0.2*w), int(0.8*w))\n    cut_y = random.randint(int(0.2*h), int(0.8*h))\n    \n    min_area_value = 0\n    min_w_value = 25\n    min_h_value = 25\n    \n    img_mosaic = np.zeros(img1.shape, dtype = img1.dtype)\n    \n    # top-left\n    img_mosaic[:cut_y, :cut_x] = img1[:cut_y, :cut_x]    \n    labels_1 = label1.copy()\n    labels_1[:, [1, 3]] = labels_1[:, [1, 3]].clip(min=1, max=cut_x)\n    labels_1[:, [2, 4]] = labels_1[:, [2, 4]].clip(min=1, max=cut_y)\n    \n    labels_1 = labels_1.astype(np.int32)\n    labels_1 = labels_1[np.where((labels_1[:,3]-labels_1[:,1])*(labels_1[:,4]-labels_1[:,2]) > min_area_value)]\n    \n    # bottom-right\n    img_mosaic[cut_y:h, cut_x:w] = img2[cut_y:h, cut_x:w]   \n    labels_2 = label2.copy()\n    labels_2[:, [1, 3]] = labels_2[:, [1, 3]].clip(min=cut_x, max=w)\n    labels_2[:, [2, 4]] = labels_2[:, [2, 4]].clip(min=cut_y, max=h)\n    \n    labels_2 = labels_2.astype(np.int32)\n    labels_2 = labels_2[np.where((labels_2[:,3]-labels_2[:,1])*(labels_2[:,4]-labels_2[:,2]) > min_area_value)]\n    \n    labels_1 = np.append(labels_1, labels_2, axis = 0)\n           \n    # top-right\n    img_mosaic[:cut_y, cut_x:w] = img3[:cut_y, cut_x:w]   \n    labels_3 = label3.copy()\n    labels_3[:, [1, 3]] = labels_3[:, [1, 3]].clip(min=cut_x, max=w)\n    labels_3[:, [2, 4]] = labels_3[:, [2, 4]].clip(min=1, max=cut_y)\n    \n    labels_3 = labels_3.astype(np.int32)\n    labels_3 = labels_3[np.where((labels_3[:,3]-labels_3[:,1])*(labels_3[:,4]-labels_3[:,2]) > min_area_value)]\n    \n    labels_1 = np.append(labels_1, labels_3, axis = 0)   \n    \n    # bottom-left\n    img_mosaic[cut_y:h, 0:cut_x] = img4[cut_y:h, 0:cut_x]   \n    labels_4 = label4.copy()\n    labels_4[:, [1, 3]] = labels_4[:, [1, 3]].clip(min=1, max=cut_x)\n    labels_4[:, [2, 4]] = labels_4[:, [2, 4]].clip(min=cut_y, max=h)\n    \n    labels_4 = labels_4.astype(np.int32)\n    labels_4 = labels_4[np.where((labels_4[:,3]-labels_4[:,1])*(labels_4[:,4]-labels_4[:,2]) > min_area_value)]\n    \n    labels_1 = np.append(labels_1, labels_4, axis = 0)\n    \n    # find boxes need check\n    label_check1 = labels_1[np.where(labels_1[:,1] == cut_x)]\n    label_check2 = labels_1[np.where(labels_1[:,3] == cut_x)]\n    label_check3 = labels_1[np.where(labels_1[:,2] == cut_y)]\n    label_check4 = labels_1[np.where(labels_1[:,4] == cut_y)]\n    \n    label_check4 = np.append(label_check4, label_check3,axis = 0)\n    label_check4 = np.append(label_check4, label_check2,axis = 0)\n    label_check4 = np.append(label_check4, label_check1,axis = 0)\n    \n    label_check4 = np.array(list(set([tuple(t) for t in label_check4])))\n    \n    if label_check4.shape[0]:\n    \n        # find boxes no need check\n        labels_all = labels_1.view([('', labels_1.dtype)] * labels_1.shape[1])\n        label_check = label_check4.view([('', label_check4.dtype)] * label_check4.shape[1])\n        label_nocheck = np.setdiff1d(labels_all, label_check).view(labels_1.dtype).reshape(-1, labels_1.shape[1])\n\n        # filter box-w and bow-h\n        label_check4 = label_check4[np.where((label_check4[:,3]-label_check4[:,1]) > min_w_value)]\n        label_check4 = label_check4[np.where((label_check4[:,4]-label_check4[:,2]) > min_h_value)]\n        label_check4 = label_check4[np.where((label_check4[:,3]-label_check4[:,1])/(label_check4[:,4]-label_check4[:,2]) <7)]\n        label_check4 = label_check4[np.where((label_check4[:,4]-label_check4[:,2])/(label_check4[:,3]-label_check4[:,1]) <7)]\n\n        label_nocheck = np.append(label_nocheck, label_check4, axis = 0)  \n        labels_1 = label_nocheck      \n    \n    \n    return img_mosaic, labels_1\n\ndef augment_mixup(img, labels, img_r, labels_r, alpha=0.5):\n    mixup_image = alpha*img + (1-alpha) * img_r\n    # print('------------------------')\n    # print(labels.shape)\n    # print(labels_r.shape)\n    labels = np.append(labels, labels_r, axis = 0)\n    # print(labels.shape)\n    \n    # image, boxes = dataset.load_image_and_boxes(random.randint(0, dataset.image_ids.shape[0] - 1))\n    # r_image, r_boxes = dataset.load_image_and_boxes(random.randint(0, dataset.image_ids.shape[0] - 1))\n    # mixup_image = alpha*image + (1-alpha)* r_image\n    # for box in boxes.astype(int):\n        # cv2.rectangle(image,(box[0], box[1]),(box[2],  box[3]),(0, 0, 1), 3)\n        # cv2.rectangle(mixup_image,(box[0], box[1]),(box[2],  box[3]),(0, 0, 1), 3)\n        \n    # for box in r_boxes.astype(int):\n        # cv2.rectangle(r_image,(box[0], box[1]),(box[2],  box[3]),(1, 0, 0), 3)\n        # cv2.rectangle(mixup_image,(box[0], box[1]),(box[2],  box[3]),(1, 0, 0), 3)\n    return mixup_image, labels\n        \ndef random_affine(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10, border=0):\n    # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))\n    # https://medium.com/uruvideo/dataset-augmentation-with-random-homographies-a8f4b44830d4\n\n    if targets is None:  # targets = [cls, xyxy]\n        targets = []\n    height = img.shape[0] + border * 2\n    width = img.shape[1] + border * 2\n\n    # Rotation and Scale\n    R = np.eye(3)\n    a = random.uniform(-degrees, degrees)\n    # a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations\n    s = random.uniform(1 - scale, 1 + scale)\n    R[:2] = cv2.getRotationMatrix2D(angle=a, center=(img.shape[1] / 2, img.shape[0] / 2), scale=s)\n\n    # Translation\n    T = np.eye(3)\n    T[0, 2] = random.uniform(-translate, translate) * img.shape[0] + border  # x translation (pixels)\n    T[1, 2] = random.uniform(-translate, translate) * img.shape[1] + border  # y translation (pixels)\n\n    # Shear\n    S = np.eye(3)\n    S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)\n    S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)\n\n    # Combined rotation matrix\n    M = S @ T @ R  # ORDER IS IMPORTANT HERE!!\n    if (border != 0) or (M != np.eye(3)).any():  # image changed\n        img = cv2.warpAffine(img, M[:2], dsize=(width, height), flags=cv2.INTER_LINEAR, borderValue=(114, 114, 114))\n\n    # Transform label coordinates\n    n = len(targets)\n    if n:\n        # warp points\n        xy = np.ones((n * 4, 3))\n        xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2)  # x1y1, x2y2, x1y2, x2y1\n        xy = (xy @ M.T)[:, :2].reshape(n, 8)\n\n        # create new boxes\n        x = xy[:, [0, 2, 4, 6]]\n        y = xy[:, [1, 3, 5, 7]]\n        xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T\n\n        # # apply angle-based reduction of bounding boxes\n        # radians = a * math.pi / 180\n        # reduction = max(abs(math.sin(radians)), abs(math.cos(radians))) ** 0.5\n        # x = (xy[:, 2] + xy[:, 0]) / 2\n        # y = (xy[:, 3] + xy[:, 1]) / 2\n        # w = (xy[:, 2] - xy[:, 0]) * reduction\n        # h = (xy[:, 3] - xy[:, 1]) * reduction\n        # xy = np.concatenate((x - w / 2, y - h / 2, x + w / 2, y + h / 2)).reshape(4, n).T\n\n        # reject warped points outside of image\n        xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)\n        xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)\n        w = xy[:, 2] - xy[:, 0]\n        h = xy[:, 3] - xy[:, 1]\n        area = w * h\n        area0 = (targets[:, 3] - targets[:, 1]) * (targets[:, 4] - targets[:, 2])\n        ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16))  # aspect ratio\n        i = (w > 4) & (h > 4) & (area / (area0 * s + 1e-16) > 0.2) & (ar < 10)\n\n        targets = targets[i]\n        targets[:, 1:5] = xy[i]\n\n    return img, targets\n    \ndef random_rotate(img, targets=(), degrees=10, border=0):\n    if targets is None:  # targets = [cls, xyxy]\n        targets = []\n    height = img.shape[0] + border * 2\n    width = img.shape[1] + border * 2\n\n    # Rotation\n    R = np.eye(3)\n    R[:2] = cv2.getRotationMatrix2D(angle=degrees, center=(img.shape[1] / 2, img.shape[0] / 2), scale=1)\n\n    # Combined rotation matrix\n    M = R\n    if (border != 0) or (M != np.eye(3)).any(): \n        img = cv2.warpAffine(img, M[:2], dsize=(width, height), flags=cv2.INTER_LINEAR, borderValue=(0, 0, 0))\n\n    n = len(targets)\n\n    if n:\n        # warp points\n        xy = np.ones((n * 4, 3))\n        xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2)  # x1y1, x2y2, x1y2, x2y1\n        xy = (xy @ M.T)[:, :2].reshape(n, 8)\n\n        # create new boxes\n        x = xy[:, [0, 2, 4, 6]]\n        y = xy[:, [1, 3, 5, 7]]\n        xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T\n\n        # reject warped points outside of image\n        xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)\n        xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)\n        w = xy[:, 2] - xy[:, 0]\n        h = xy[:, 3] - xy[:, 1]\n        area = np.fabs( w * h )\n        area0 = np.fabs((targets[:, 3] - targets[:, 1]) * (targets[:, 4] - targets[:, 2]))\n        ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16))  # aspect ratio\n        i = (w > 4) & (h > 4) & (area / (area0 * 1 + 1e-16) > 0.2) & (ar < 10)\n        targets = targets[i]\n        targets[:, 1:5] = xy[i]\n\n    return img, targets\n\n    \ndef augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):\n    r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains\n    hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))\n    dtype = img.dtype  # uint8\n\n    x = np.arange(0, 256, dtype=np.int16)\n    lut_hue = ((x * r[0]) % 180).astype(dtype)\n    lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)\n    lut_val = np.clip(x * r[2], 0, 255).astype(dtype)\n\n    img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)\n    cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)  # no return needed\n\ndef augment_hsv2(img):\n    flag = False\n    num_count = 0\n    \n    img_original = copy.deepcopy(img)\n    \n    while num_count<50:\n        hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_RGB2HSV))\n        mean_before = np.mean(hue)\n        \n        value_hue = random.randint(-20,20)\n        value_sat = random.randint(80,95)*0.01\n        value_exp = random.randint(80,95)*0.01        \n\n        hue = (hue + value_hue).astype(img.dtype)\n        sat = np.clip(sat * value_sat, 0, 255).astype(img.dtype)\n        val = np.clip(val * value_exp, 0, 255).astype(img.dtype)\n        \n        mean_after = np.mean(hue) \n        num_count = num_count+1\n        if  int(np.fabs(mean_after -  mean_before)) in range(20) and int(mean_after)<110 and int(mean_after)>15:\n            img_hsv = cv2.merge((hue, sat, val)).astype(img.dtype)\n            cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB, dst=img)\n            flag = True\n    \n    if flag:\n        return img\n    else:\n        return img_original\n\ndef augment_flip(img, labels):\n    r = np.random.randint(0, 2) -1\n    cv2.flip(img, r, img)\n    img_h = img.shape[0]\n    img_w = img.shape[1]\n\n    if r == 0:\n        labels[:,2] = img_h - labels[:,2]\n        labels[:,4] = img_h - labels[:,4]\n        labels = labels[:,[0,1,4,3,2]]\n    elif r == -1:\n        labels[:,1] = img_w - labels[:,1]\n        labels[:,3] = img_w - labels[:,3]\n        labels[:,2] = img_h - labels[:,2]\n        labels[:,4] = img_h - labels[:,4]\n        labels = labels[:,[0,3,4,1,2]]        \n    \n    return img, labels\n\ndef augment_light_contrast(img):\n    # dst = alpha * img + beta * blank\n    alpha = np.random.uniform(0, 3, 1)*0.1 + 1\n    beta = random.randint(-15,15)\n    blank = np.zeros(img.shape, img.dtype)\n    dst = cv2.addWeighted(img, alpha, blank, 1-alpha, beta)  \n    dst = np.clip(dst, 0, 255).astype(img.dtype)\n    return dst\n\ndef augment_illumination(img, scale, normallization):\n    height, width = img.shape[:2]\n    # IMAGE_WIDTH = 512\n    # IMAGE_HEIGHT = 392\n    center_x = width/2\n    center_y = height/2\n\n    R = np.sqrt(center_x ** 2 + center_y ** 2) * scale\n\n    Gauss_map = np.zeros((height, width))\n\n    for i in range(height):\n        for j in range(width):\n            dis = np.sqrt((i - center_y) ** 2 + (j - center_x) ** 2)\n            Gauss_map[i, j] = np.exp(-0.5 * dis / R)\n\n    mask_x = repmat(center_x, height, width)\n    mask_y = repmat(center_y, height, width)\n\n    x1 = np.arange(width)\n    x_map = repmat(x1, height, 1)\n\n    y1 = np.arange(height)\n    y_map = repmat(y1, width, 1)\n    y_map = np.transpose(y_map)\n\n    Gauss_map = np.sqrt((x_map - mask_x) ** 2 + (y_map - mask_y) ** 2)\n\n    Gauss_map = np.exp(-0.5 * Gauss_map / R)\n\n    illumination = np.zeros([height, width, 3], np.float32)\n\n    if normallization:\n        gaussian = abs(Gauss_map)\n        max_gaussian = np.max(gaussian)\n        min_gaussian = np.min(gaussian)\n        gaussian = (gaussian - min_gaussian) / (max_gaussian - min_gaussian)\n\n        illumination[:, :, 0] = gaussian\n        illumination[:, :, 1] = gaussian\n        illumination[:, :, 2] = gaussian\n    else:\n        illumination[:, :, 0] = Gauss_map\n        illumination[:, :, 1] = Gauss_map\n        illumination[:, :, 2] = Gauss_map\n\n    illumination_img = img * illumination\n\n    MAX = 255\n    inds = np.where(\n        (illumination_img[:, :, 0] > MAX) &\n        (illumination_img[:, :, 1] > MAX) &\n        (illumination_img[:, :, 2] > MAX))[0]\n    illumination_img[inds, :] = 255\n\n    MIN = 0\n    inds = np.where(\n        (illumination_img[:, :, 0] < MIN) &\n        (illumination_img[:, :, 1] < MIN) &\n        (illumination_img[:, :, 2] < MIN))[0]\n    illumination_img[inds, :] = 0\n\n    illumination_img = np.uint8(illumination_img)\n    # img = cv2.equalizeHist(img)\n\n    return illumination_img\n\n\n\n#######################################################################################\n\n\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')    \nprint('platform: ', device)\n\nimport random, math\n# Albumentations\ndef get_train_transform():\n    return A.Compose([\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']})\n\ndef get_test_transform():\n    return A.Compose([\n        # A.Resize(512, 512),\n        ToTensorV2(p=1.0)\n    ])\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_dataset = WheatDataset(train_df, DIR_TRAIN, \"train\", get_train_transform())\nvalid_dataset = WheatDataset(valid_df, DIR_TRAIN, \"val\", get_valid_transform())\ntest_dataset = WheatTestDataset(test_df, DIR_TEST, get_test_transform())\n\n#train_dataset = WheatDataset(train_df, DIR_TRAIN, get_augumentation(phase='train'))\n\n# split the dataset in train and test set\n# indices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=16,\n    shuffle=True,\n    num_workers=8,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=16,\n    shuffle=False,\n    num_workers=8,\n    collate_fn=collate_fn\n)\n\ntest_data_loader = DataLoader(\n    test_dataset,\n    batch_size=16,\n    shuffle=False,\n    num_workers=8,\n    drop_last=False,\n    collate_fn=collate_fn\n)\n\nprint('---------------network configuration...--------------------')\n########################## train faster rcnn with a specified backbone ###########\n# backbone = torchvision.models.mobilenet_v2(pretrained=True).features\n# backbone.out_channels = 1280\n# anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),))\n# roi_pooler = torchvision.ops.MultiScaleRoIAlign(feature_names=[0], output_size = 7, sampling_ratio=2)\n# model = FasterRCNN(backbone, num_classes=2, rpn_anchor_generator=anchor_generator, box_roi_pool = roi_pooler)\n\n########################## finetune from a model pretrained on COCO ###########\n# load a model pre-trained on COCO\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)\n\n#replace the classifier with a new one\n# num_classes which is user-defined\nnum_classes = 2  # 1 class (wheat) + background\n\n# get number of input features for the classifier\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n\n# replace the pre-trained head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\nWEIGHTS_FILE = os.path.join(DIR_INPUT, \"fasterrcnn/fasterrcnn_resnet50_fpn_best.pth\")\nmodel.load_state_dict(torch.load(WEIGHTS_FILE, map_location=device))\n\n    \nmodel.to(device)\nparams = [p for p in model.parameters() if p.requires_grad]\n# optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\n\noptimizer = torch.optim.SGD(params, lr=0.01, momentum=0.9, weight_decay=0.0005)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)\n#lr_scheduler = None\n\n\nprint('---------------begin Training...--------------------')\nnum_epochs = 100 #Increase it for better results\ndisplay_interval = 50\n\nloss_hist = Averager()\nval_loss_hist = Averager()\nitr = 1\nleast_loss = float('inf')\nfor epoch in range(num_epochs):\n    loss_hist.reset()\n    val_loss_hist.reset()\n    \n    for images, targets in train_data_loader:\n        # print('**************************')\n        # for img in images:\n            # print(\"type(img): \", type(img))\n            # print(img)\n        images = list(image.to(device) for image in images)\n        # for t in targets:\n            # print(t)\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n        \n        loss_dict = model(images, targets)\n\n        losses = sum(loss for loss in loss_dict.values())\n        loss_value = losses.item()\n\n        loss_hist.send(loss_value)\n\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n\n        if itr % display_interval == 0:\n            print(\"Iteration #%s loss: %s\" %(itr, loss_value))\n\n        itr += 1\n    \n    #Validation Step\n    for images, targets in valid_data_loader:\n        \n        images = list(image.to(device) for image in images)\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n        val_loss_dict = model(images, targets)\n\n        val_losses = sum(loss for loss in val_loss_dict.values())\n        val_loss_value = val_losses.item()\n\n        val_loss_hist.send(val_loss_value)\n\n    \n    # update the learning rate\n    if loss_hist.value<least_loss:\n        least_loss = loss_hist.value  #average loss of all iters once epoch\n        ltrain=int(least_loss*1000)/1000\n        print('save model epoch= %s, loss = %s' % (epoch, ltrain))\n        # torch.save(model.state_dict(), 'fasterrcnn_custom_test_ep%s_loss%s.pth' % (epoch, ltrain))\n        torch.save(model.state_dict(), 'best_weights.pth')        \n    else:\n        if lr_scheduler is not None:\n            lr_scheduler.step()\n            \n    #if val_loss_hist.value<least_loss:\n    #    least_loss = val_loss_hist.value  #average loss of all iters once epoch\n    #    lval=int(least_loss*1000)/1000\n    #    torch.save(model.state_dict(), 'fasterrcnn_custom_test_ep%s_loss%s.pth' % (epoch, lval))\n    #    # torch.save(model.state_dict(), 'best_weights.pth')        \n    #else:\n    #    if lr_scheduler is not None:\n    #        lr_scheduler.step()\n            \n    print(\"Epoch #%s train_loss: %s val_loss: %s\" % (epoch, loss_hist.value, val_loss_hist.value))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"begin Inference and prepare submission","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"############################## Inference #################################\nprint('---------------begin inference...--------------------')\n\ndef format_prediction_string(boxes, scores):\n    pred_strings = []\n    for j in zip(scores, boxes):\n        pred_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))\n\n    return \" \".join(pred_strings)\n\nWEIGHTS_FILE = 'best_weights.pth'\nmodel.load_state_dict(torch.load(WEIGHTS_FILE, map_location=device))\nmodel.eval()\n\ndetection_threshold = 0.5\nresults = []\n\nfor images, image_ids in test_data_loader:\n\n    images = list(image.to(device) for image in images)\n    outputs = model(images)\n\n    for i, image in enumerate(images):\n        \n        boxes = outputs[i]['boxes'].data.cpu().numpy()\n        scores = outputs[i]['scores'].data.cpu().numpy()\n        \n        boxes = boxes[scores >= detection_threshold].astype(np.int32)\n        scores = scores[scores >= detection_threshold]\n        image_id = image_ids[i]\n        \n        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n        \n        result = {\n            'image_id': image_id,\n            'PredictionString': format_prediction_string(boxes, scores)\n        }\n        \n        # im_ori = cv2.imread( os.path.join(DIR_TEST, image_id+\".jpg\"))\n        # for box in boxes:\n            # cv2.rectangle(im_ori, (box[0], box[1]), (box[2]+ box[0], box[3]+box[1]), (220, 0, 0), 2)\n        \n        # if not os.path.exists(\"vis_result\"):\n            # os.makedirs(\"vis_result\")\n        # cv2.imwrite(os.path.join(\"vis_result\", image_id+\"_det.jpg\"), im_ori)\n        \n        results.append(result)\n    \n    \n    \ntest_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.head()\ntest_df.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"All is finished!\")","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}