{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install ../input/weightedboxesfusion/Weighted-Boxes-Fusion-master/ > /dev/null\n!pip install --no-deps '../input/timm-0130/timm-0.1.30-py3-none-any.whl' > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\nimport os\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom ensemble_boxes import *\nimport sys\nsys.path.append('../input/detrmodel/detr')\nfrom models.resnestbackbone import ResnestBackBone\nfrom models.matcher import HungarianMatcher\nfrom models.detr import SetCriterion\n\nfrom torch.utils.data import Dataset,DataLoader, ConcatDataset\n\nfrom util.misc import (NestedTensor, nested_tensor_from_tensor_list)\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nfrom tqdm import tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!mkdir -p /root/.cache/torch/hub/\n!cp -r ../input/detrmodel/detr /root/.cache/torch/hub/facebookresearch_detr_master","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return A.Compose(\n        [  \n          A.OneOf([\n                A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2,\n                                     val_shift_limit=0.2, p=0.3), \n                A.RandomBrightnessContrast(brightness_limit=0.2,  \n                                           contrast_limit=0.2, p=0.3),\n                A.RGBShift(r_shift_limit=20/255, g_shift_limit=20/255, b_shift_limit=20/255,p=0.3),\n            ], p=0.2),\n            A.OneOf([\n                A.RandomGamma(gamma_limit=(80, 120), p=0.3),\n                A.Blur(p=0.05),\n                A.GaussNoise(var_limit=(0.05, 0.1), mean=0, p=0.3),\n                A.ToGray(p=0.05)], p=0.2),\n\n            A.OneOf([\n                A.HorizontalFlip(p=0.75), \n                A.VerticalFlip(p=0.75),  \n                A.Transpose(p=0.75),                \n                A.RandomRotate90(p=0.75)\n                ], p=1),         \n            #  A.RandomFog(fog_coef_lower=0.1, fog_coef_upper=0.4, p=0.1), \n             A.Resize(height=1024, width=1024, p=1),\n            #  A.Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=0, p=0.3),\n             ToTensorV2(p=5.0),\n             ],\n             \n        p=1.0, bbox_params=A.BboxParams(format='coco',min_area=0, min_visibility=0,label_fields=['labels'])\n    )\ndef collate_fn(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n    def __init__(self,image_ids,dataframe,transforms=None, dir_train ='/content/datasets/ori/train'):\n        self.image_ids = image_ids\n        self.df = dataframe\n        self.transforms = transforms\n        self.dir = dir_train\n        \n        \n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n    \n    def __getitem__(self,index):\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n        image = cv2.imread(f'{self.dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        # DETR takes in data in coco format \n        boxes = records[['x', 'y', 'w', 'h']].values\n        \n        #Area of bb\n        area = boxes[:,2]*boxes[:,3]\n        area = torch.as_tensor(area, dtype=torch.float32)\n        \n        # AS pointed out by PRVI It works better if the main class is labelled as zero\n        labels =  np.zeros(len(boxes), dtype=np.int32)\n\n        \n        if self.transforms:\n            sample = {\n                'image': image,\n                'bboxes': boxes,\n                'labels': labels\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n            boxes = sample['bboxes']\n            labels = sample['labels']\n            \n            \n        #Normalizing BBOXES\n            \n        _,h,w = image.shape\n        boxes = A.augmentations.bbox_utils.normalize_bboxes(sample['bboxes'],rows=h,cols=w)\n        target = {}\n        target['boxes'] = torch.as_tensor(boxes,dtype=torch.float32)\n        target['labels'] = torch.as_tensor(labels,dtype=torch.long)\n        target['image_id'] = torch.tensor([index])\n        target['area'] = area\n        \n        return image, target, image_id","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DETRModel(nn.Module):\n    def __init__(self,num_classes,num_queries):\n        super(DETRModel,self).__init__()\n        self.num_classes = num_classes\n        self.num_queries = num_queries\n        \n        self.model = torch.hub.load('facebookresearch/detr', 'detr_resnet101', pretrained=False)\n        self.in_features = self.model.class_embed.in_features\n        \n        hidden_dim = self.model.transformer.d_model\n\n        self.model.class_embed = nn.Linear(in_features=self.in_features,out_features=self.num_classes)\n        self.model.num_queries = self.num_queries\n        self.model.backbone = ResnestBackBone()\n        self.model.input_proj = nn.Conv2d(1024, hidden_dim, kernel_size=1,groups=4)\n        self.model.query_embed = nn.Embedding(num_queries, hidden_dim)\n\n\n    def forward(self,images):\n        if isinstance(images, (list, torch.Tensor)):\n            images = nested_tensor_from_tensor_list(images)\n        return self.model(images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def tta_image(image, device):\n    image_src = torch.tensor(image.copy(),dtype=torch.float,device = device).permute([2, 0, 1])/255\n    image_hf = torch.tensor(cv2.flip(image.copy(), 1) ,dtype=torch.float,device = device).permute([2, 0, 1])/255    \n    image_vf = torch.tensor(cv2.flip(image.copy(), 0) ,dtype=torch.float,device = device).permute([2, 0, 1])/255  \n    image_hf_vf = torch.tensor(cv2.flip(image.copy(), -1) ,dtype=torch.float,device = device).permute([2, 0, 1])/255\n    return [image_src,image_hf,image_vf,image_hf_vf]\n\ndef tta_reverse_boxes(boxes):\n    boxes = (boxes).detach().cpu().numpy().astype(np.float32)\n    boxes[1,:,0] = 1 - boxes[1,:,0] - boxes[1,:,2]\n    boxes[2,:,1] = 1 - boxes[2,:,1] - boxes[2,:,3]\n    boxes[3,:,0] = 1 - boxes[3,:,0] - boxes[3,:,2]\n    boxes[3,:,1] = 1 - boxes[3,:,1] - boxes[3,:,3]\n    return boxes\n\ndef boxes_to_x1x2y1y2(boxes):\n    boxes[:,:,2] = boxes[:,:,0] + boxes[:,:,2]\n    boxes[:,:,3] = boxes[:,:,1] + boxes[:,:,3]\n    return np.clip(boxes,0,1)\n\ndef boxes_to_xywh(boxes):\n    boxes[:,2] =  boxes[:,2] - boxes[:,0]\n    boxes[:,3] =  boxes[:,3] - boxes[:,1]\n    return boxes","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_fn(data_loader,model,criterion,optimizer,device,scheduler,epoch):\n    model.train()\n    model.model.backbone.freeze_bn()\n    criterion.train()\n    \n    summary_loss = AverageMeter()\n    \n    tk0 = tqdm(data_loader, total=len(data_loader))\n    \n    for step, (images, targets, image_ids) in enumerate(tk0):\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        output = model(images)\n        for k in output.keys():\n            output[k] = output[k].float()\n\n        loss_dict = criterion(output, targets)\n        weight_dict = criterion.weight_dict\n        \n        losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)\n        optimizer.zero_grad()\n\n#         losses.backward()\n        with amp.scale_loss(losses, optimizer) as scaled_loss:\n            scaled_loss.backward()\n\n        optimizer.step()\n        if scheduler is not None:\n            scheduler.step()\n        \n        summary_loss.update(losses.item(),BATCH_SIZE)\n        tk0.set_postfix(loss=summary_loss.avg, loss_ce=loss_dict['loss_ce'].item(), loss_bbox=loss_dict['loss_bbox'].item(), loss_giou=loss_dict['loss_giou'].item(), loss_card=loss_dict['cardinality_error'].item(), acc=100-loss_dict['class_error'].item())\n        \n    return summary_loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_pl(epochs, model):\n    \n    df_train = pd.read_csv('../input/global-wheat-detection/train.csv')\n    bboxs = np.stack(df_train['bbox'].apply(lambda x: np.fromstring(x[1:-1], sep=',')))\n    for i, column in enumerate(['x', 'y', 'w', 'h']):\n        df_train[column] = bboxs[:,i]\n    df_train.drop(columns=['bbox'], inplace=True)\n    \n#     df_pl = pd.read_csv('detr_pl.csv')\n#     bboxs = np.stack(df_pl['bbox'].apply(lambda x: np.fromstring(x[1:-1], sep=',')))\n#     for i, column in enumerate(['x', 'y', 'w', 'h']):\n#         df_pl[column] = bboxs[:,i]\n#     df_pl.drop(columns=['bbox'], inplace=True)\n    \n    train_dataset_ori = WheatDataset(\n    image_ids=df_train['image_id'].unique(),\n    dataframe=df_train,\n    transforms=get_train_transforms(),\n    dir_train ='../input/global-wheat-detection/train'\n\n    )\n    \n#     train_dataset_pl = WheatDataset(\n#     image_ids=df_pl['image_id'].unique(),\n#     dataframe=df_pl,\n#     transforms=get_train_transforms(),\n#     dir_train ='../input/global-wheat-detection/test'\n#     )\n    \n    \n#     train_dataset = ConcatDataset([train_dataset_ori, train_dataset_pl])\n    train_dataset = train_dataset_ori\n\n\n    train_data_loader = DataLoader(\n    train_dataset,\n    batch_size=1,\n    shuffle=True,\n    num_workers=4,\n    collate_fn=collate_fn\n    )\n\n    \n    weight_dict = {'loss_ce': 1, 'loss_bbox': 1 , 'loss_giou': 1}\n    matcher = HungarianMatcher(cost_class=weight_dict['loss_ce'], cost_bbox=weight_dict['loss_bbox'], cost_giou=weight_dict['loss_giou'], multiple_match=1)\n    device = torch.device('cuda')\n    model = model.to(device)\n    criterion = SetCriterion(1, matcher, weight_dict, eos_coef = 0.5, losses=['labels', 'boxes', 'cardinality'])\n    criterion = criterion.to(device)\n    \n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)\n    model, optimizer = amp.initialize(model, optimizer, opt_level='O1')\n    scheduler = None\n    for epoch in range(epochs):\n        train_loss = train_fn(train_data_loader, model,criterion, optimizer,device,scheduler=scheduler,epoch=epoch)\n        torch.save(model.state_dict(), './detr_pl.pt')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def inference(model, device, source = '../input/global-wheat-detection/test/', view_visual=False, pl=False):\n    image_file_names =  os.listdir(source)\n    results = []\n    for file_name in image_file_names:\n        image_id = file_name.split('.')[0]\n        image = cv2.imread('%s/%s.jpg'%(source,image_id), cv2.IMREAD_COLOR)\n        image = cv2.resize(image,(1024,1024))\n        image_raw = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image = tta_image(image_raw, device)\n        output = model(image)\n        del image\n\n        boxes = tta_reverse_boxes(output['pred_boxes'])\n        prob   = output['pred_logits'].softmax(dim=-1).detach().cpu().numpy()[:,:,0]\n        boxes, prob, _ = weighted_boxes_fusion(boxes_to_x1x2y1y2(boxes), prob, np.zeros_like(prob),iou_thr=0.4, skip_box_thr=0.1)\n        boxes = (boxes_to_xywh(boxes)*1024).astype(np.int32)\n        \n\n        if view_visual:\n            for box,p in zip(boxes,prob):\n                if p >0.5:\n                    color = (255,0,0) #if p>0.5 else (0,0,0)\n                    font = cv2.FONT_HERSHEY_SIMPLEX\n                    cv2.putText(image_raw,str(p),(box[0], box[1]),font,0.5,color,1)\n                    cv2.rectangle(image_raw,(box[0], box[1]),(box[2]+box[0], box[3]+box[1]),color, 1)\n            cv2.imwrite(f'./{image_id}.jpg',cv2.cvtColor(image_raw, cv2.COLOR_RGB2BGR))\n            \n        if not pl:\n            prediction_strings = []\n            for p,b in zip(prob,boxes):\n                if p > 0.5:\n                    box = b.astype(np.int32)\n                    prediction_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(p, box[0],box[1],box[2],box[3]))\n\n            r = {\n                'image_id':image_id,\n                'PredictionString':\" \".join(prediction_strings)\n                }\n            results.append(r)\n        else:\n            for p,b in zip(prob,boxes):\n                if p > 0.5:\n                    box = b.astype(np.int32)\n                    r = {\n                        'image_id':image_id,\n                        'width': '1024',\n                        'height' : '1024',\n                        'bbox' : f'[{box[0]},{box[1]},{box[2]},{box[3]}]',\n                        'source': 'detr_pl'\n                        }\n                    results.append(r)\n    return results","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_model(path = '../input/detrmodel/detr_size1024_q512_val40.pt'):\n    model = DETRModel(2, 512)\n    model.load_state_dict(torch.load(path,map_location=device))\n    model.eval()\n    model = model.to(device)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if len(os.listdir('../input/global-wheat-detection/test/'))<11:\n    committing=True\nelse:\n    committing=False\n    \nmodel = get_model()\nwith torch.no_grad():\n    results = inference(model,device,view_visual=committing)\ntest_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('submission.csv', index=False)\ntest_df.head()","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}