{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Referece:\n#### [Face Mask Detection Fast](https://www.kaggle.com/mekhdigakhramanian/face-mask-detection-fast)\n#### [DETR model](https://www.kaggle.com/tanulsingh077/end-to-end-object-detection-with-transformers-detr)","metadata":{}},{"cell_type":"code","source":"!git clone https://github.com/facebookresearch/detr.git","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:45.769188Z","iopub.execute_input":"2022-03-10T02:44:45.769730Z","iopub.status.idle":"2022-03-10T02:44:48.433035Z","shell.execute_reply.started":"2022-03-10T02:44:45.769637Z","shell.execute_reply":"2022-03-10T02:44:48.430651Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport random\n\nfrom tqdm import tqdm\nimport xml.etree.ElementTree as ET\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\nimport sys\nsys.path.append('./detr/')\n\nfrom detr.models.matcher import HungarianMatcher\nfrom detr.models.detr import SetCriterion\n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport cv2\nimport matplotlib.patches as patches\nimport matplotlib.text as text\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.model_selection import StratifiedKFold","metadata":{"execution":{"iopub.status.busy":"2022-03-10T03:11:01.619540Z","iopub.execute_input":"2022-03-10T03:11:01.619806Z","iopub.status.idle":"2022-03-10T03:11:01.627425Z","shell.execute_reply.started":"2022-03-10T03:11:01.619776Z","shell.execute_reply":"2022-03-10T03:11:01.626324Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed = 42\nnum_classes = 4\nnum_queries = 20\nnull_class_coef = 0.2\nBATCH_SIZE = 32\nLR = 2e-5\nEPOCHS = 50\nANNOT_DIR = '../input/face-mask-detection/annotations'\ndevice = torch.device('cuda')\n\ndef get_objects(xml_file):\n    annotation = ET.parse(xml_file)\n    root = annotation.getroot()\n    \n    name = os.path.basename(xml_file).replace('.xml','')\n    size = root.find('size')\n    \n    objects = []\n    for obj in root.findall('object'):\n        bbox = obj.find('bndbox')\n        new_object = {}\n        new_object['image_id'] = name\n        new_object['labels'] = obj.find('name').text\n        new_object['width'] = int(size.find('width').text)\n        new_object['height'] = int(size.find('height').text)\n        new_object['x'] = int(bbox.find('xmin').text)\n        new_object['y'] = int(bbox.find('ymin').text)\n        new_object['w'] = int(bbox.find('xmax').text)-int(bbox.find('xmin').text)\n        new_object['h'] = int(bbox.find('ymax').text)-int(bbox.find('ymin').text)\n        objects.append(new_object)\n    return objects","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:52.306667Z","iopub.execute_input":"2022-03-10T02:44:52.306931Z","iopub.status.idle":"2022-03-10T02:44:52.318625Z","shell.execute_reply.started":"2022-03-10T02:44:52.306885Z","shell.execute_reply":"2022-03-10T02:44:52.317989Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\nseed_everything(seed)","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:52.321063Z","iopub.execute_input":"2022-03-10T02:44:52.321371Z","iopub.status.idle":"2022-03-10T02:44:52.329620Z","shell.execute_reply.started":"2022-03-10T02:44:52.321310Z","shell.execute_reply":"2022-03-10T02:44:52.328921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"annots = []\nfor xml in os.listdir(ANNOT_DIR):\n    annots += get_objects(os.path.join(ANNOT_DIR,xml))\ndf = pd.DataFrame(annots)\n\n# Adjust w, h\nfor i in df.index:\n    exceed_w =  df.iloc[i].x + df.iloc[i].w - df.iloc[i].width\n    exceed_h =  df.iloc[i].y + df.iloc[i].h - df.iloc[i].height\n    if exceed_w > 0:\n        df.loc[df.index == i,'w'] -= exceed_w\n    if exceed_h > 0:\n        df.loc[df.index == i,'h'] -= exceed_h\n\nencoder = LabelEncoder()\ndf.labels = encoder.fit_transform(df.labels)\n\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:52.330852Z","iopub.execute_input":"2022-03-10T02:44:52.331869Z","iopub.status.idle":"2022-03-10T02:44:58.528169Z","shell.execute_reply.started":"2022-03-10T02:44:52.331813Z","shell.execute_reply":"2022-03-10T02:44:58.527497Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"g = sns.countplot(x='labels',data=df)\ng.set_xticklabels(encoder.classes_);","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:58.529377Z","iopub.execute_input":"2022-03-10T02:44:58.529706Z","iopub.status.idle":"2022-03-10T02:44:58.726817Z","shell.execute_reply.started":"2022-03-10T02:44:58.529668Z","shell.execute_reply":"2022-03-10T02:44:58.726059Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_split = df[['image_id']].copy()\n# df_split = pd.concat([df_split,pd.get_dummies(df.labels).astype('int32')],axis=1)\ndf_split['bbox_count'] = 1\ndf_split = df_split.groupby('image_id').sum()\n\ng = plt.figure(figsize=(15,4))\ng = sns.countplot(x='bbox_count',data=df_split)","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:58.729326Z","iopub.execute_input":"2022-03-10T02:44:58.729616Z","iopub.status.idle":"2022-03-10T02:44:59.112639Z","shell.execute_reply.started":"2022-03-10T02:44:58.729576Z","shell.execute_reply":"2022-03-10T02:44:59.111949Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_split","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.113944Z","iopub.execute_input":"2022-03-10T02:44:59.114205Z","iopub.status.idle":"2022-03-10T02:44:59.124633Z","shell.execute_reply.started":"2022-03-10T02:44:59.114162Z","shell.execute_reply":"2022-03-10T02:44:59.122894Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_split = df_split[:int(len(df_split)*.9)]\nval_split = df_split[int(len(df_split)*.9):]","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.126110Z","iopub.execute_input":"2022-03-10T02:44:59.127000Z","iopub.status.idle":"2022-03-10T02:44:59.131667Z","shell.execute_reply.started":"2022-03-10T02:44:59.126962Z","shell.execute_reply":"2022-03-10T02:44:59.130945Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axs = plt.subplots(1,2,figsize=(13,4))\nsns.countplot(x='labels',data=df[df['image_id'].isin(train_split.index)],ax=axs[0])\naxs[0].set_title('Labels distribution in training set')\nsns.countplot(x='labels',data=df[df['image_id'].isin(val_split.index)],ax=axs[1])\naxs[1].set_title('Labels distribution in validation set')","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.134361Z","iopub.execute_input":"2022-03-10T02:44:59.134741Z","iopub.status.idle":"2022-03-10T02:44:59.408365Z","shell.execute_reply.started":"2022-03-10T02:44:59.134704Z","shell.execute_reply":"2022-03-10T02:44:59.407699Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_transform():\n    return A.Compose([\n        A.HorizontalFlip(),\n        A.RandomRotate90(),\n        A.RandomBrightnessContrast(),\n        A.Resize(300, 300),\n        ToTensorV2(),\n    ], bbox_params=A.BboxParams(format='coco',label_fields=['labels']))\n\ndef valid_trainsform():\n    return A.Compose([\n        A.Resize(300,300),\n        ToTensorV2(),\n    ], bbox_params=A.BboxParams(format='coco',label_fields=['labels']))","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.409393Z","iopub.execute_input":"2022-03-10T02:44:59.409780Z","iopub.status.idle":"2022-03-10T02:44:59.416909Z","shell.execute_reply.started":"2022-03-10T02:44:59.409744Z","shell.execute_reply":"2022-03-10T02:44:59.416151Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IMG_DIR = '../input/face-mask-detection/images'\n\nclass MaskDataset(Dataset):\n    def __init__(self, image_ids, dataframe, transforms=None):\n        self.image_ids = image_ids\n        self.df = dataframe\n        self.transforms = transforms\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n    def __getitem__(self, index):\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id']==image_id]\n        \n        image = cv2.imread(f'{IMG_DIR}/{image_id}.png', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        \n        boxes = records[['x', 'y', 'w', 'h']].values\n        area = boxes[:,2]*boxes[:,3]\n        area = torch.as_tensor(area, dtype=torch.float32)\n        \n        labels =  records['labels'].values\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        _, h, w = image.shape\n        boxes = A.augmentations.bbox_utils.normalize_bboxes(sample['bboxes'],rows=h,cols=w)  \n        \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","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.419235Z","iopub.execute_input":"2022-03-10T02:44:59.419506Z","iopub.status.idle":"2022-03-10T02:44:59.432959Z","shell.execute_reply.started":"2022-03-10T02:44:59.419474Z","shell.execute_reply":"2022-03-10T02:44:59.432308Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.434286Z","iopub.execute_input":"2022-03-10T02:44:59.434704Z","iopub.status.idle":"2022-03-10T02:44:59.444576Z","shell.execute_reply.started":"2022-03-10T02:44:59.434669Z","shell.execute_reply":"2022-03-10T02:44:59.443814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = MaskDataset(image_ids=train_split.index.values,\n                           dataframe=df,\n                           transforms=train_transform())\ntrain_data_loader = DataLoader(train_dataset,\n                             batch_size=BATCH_SIZE,\n                             shuffle=False,\n                             num_workers=2,\n                             collate_fn=collate_fn)\nval_dataset = MaskDataset(image_ids=val_split.index.values,\n                           dataframe=df,\n                           transforms=valid_trainsform())\nval_data_loader = DataLoader(val_dataset,\n                             batch_size=BATCH_SIZE,\n                             shuffle=False,\n                             num_workers=2,\n                             collate_fn=collate_fn)","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.447648Z","iopub.execute_input":"2022-03-10T02:44:59.447831Z","iopub.status.idle":"2022-03-10T02:44:59.454804Z","shell.execute_reply.started":"2022-03-10T02:44:59.447809Z","shell.execute_reply":"2022-03-10T02:44:59.453873Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"matcher = HungarianMatcher()\nweight_dict = {'loss_ce': 1, 'loss_bbox': 1 , 'loss_giou': 1}\nlosses = ['labels', 'boxes', 'cardinality']","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.455978Z","iopub.execute_input":"2022-03-10T02:44:59.456358Z","iopub.status.idle":"2022-03-10T02:44:59.466340Z","shell.execute_reply.started":"2022-03-10T02:44:59.456322Z","shell.execute_reply":"2022-03-10T02:44:59.465612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AverageMeter(object):\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","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.467781Z","iopub.execute_input":"2022-03-10T02:44:59.468088Z","iopub.status.idle":"2022-03-10T02:44:59.476180Z","shell.execute_reply.started":"2022-03-10T02:44:59.468045Z","shell.execute_reply":"2022-03-10T02:44:59.475449Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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        self.model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)\n        self.in_features = self.model.class_embed.in_features\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    def forward(self, images):\n        return self.model(images)","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:44:59.698108Z","iopub.execute_input":"2022-03-10T02:44:59.698354Z","iopub.status.idle":"2022-03-10T02:44:59.705480Z","shell.execute_reply.started":"2022-03-10T02:44:59.698326Z","shell.execute_reply":"2022-03-10T02:44:59.704593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_fn(data_loader, model, criterion, optimizer, device, epoch):\n    model.train()\n    criterion.train()\n    \n    summary_loss = AverageMeter()\n    \n    for step, (images, targets, image_ids) in enumerate(data_loader):\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        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\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n        \n        summary_loss.update(losses.item(), BATCH_SIZE)\n    return summary_loss\n\ndef eval_fn(data_loader, model, criterion, device, epoch):\n    model.eval()\n    criterion.eval()\n    \n    summary_loss = AverageMeter()\n    \n    with torch.no_grad():\n        for step, (images, targets, image_ids) in enumerate(data_loader):\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            loss_dict = criterion(output, targets)\n            weight_dict = criterion.weight_dict\n            losses = sum(loss_dict[k]*weight_dict[k] for k in loss_dict.keys() if k in weight_dict)\n            \n            summary_loss.update(losses.item(), BATCH_SIZE)\n    return summary_loss","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:45:00.759567Z","iopub.execute_input":"2022-03-10T02:45:00.760068Z","iopub.status.idle":"2022-03-10T02:45:00.771325Z","shell.execute_reply.started":"2022-03-10T02:45:00.760031Z","shell.execute_reply":"2022-03-10T02:45:00.770622Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.environ['WANDB_CONSOLE'] = 'off'\n\ndef run():\n    model = DETRModel(num_classes=num_classes, num_queries=num_queries)\n    model = model.to(device)\n    criterion = SetCriterion(num_classes-1, matcher, weight_dict, eos_coef=null_class_coef, losses=losses)\n    criterion = criterion.to(device)\n    optimizer = torch.optim.AdamW(model.parameters(), lr=LR)\n    best_loss = 10**5\n    for epoch in range(1,EPOCHS+1):\n        train_loss = train_fn(train_data_loader, model, criterion, optimizer, device, epoch=epoch)\n        valid_loss = eval_fn(val_data_loader, model, criterion, device, epoch)\n        if valid_loss.avg < best_loss:\n            best_loss = valid_loss.avg\n            torch.save(model.state_dict(), f'detr_best.pth')\n        if epoch%10 == 0:\n            print(f'Epoch {epoch+0:03}: | Train Loss: {train_loss.avg:.5f} | Val Loss: {valid_loss.avg:.5f}')","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:45:18.154101Z","iopub.execute_input":"2022-03-10T02:45:18.154394Z","iopub.status.idle":"2022-03-10T02:45:18.162550Z","shell.execute_reply.started":"2022-03-10T02:45:18.154355Z","shell.execute_reply":"2022-03-10T02:45:18.161423Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = run()","metadata":{"execution":{"iopub.status.busy":"2022-03-10T02:45:19.598171Z","iopub.execute_input":"2022-03-10T02:45:19.598909Z","iopub.status.idle":"2022-03-10T03:02:53.664168Z","shell.execute_reply.started":"2022-03-10T02:45:19.598873Z","shell.execute_reply":"2022-03-10T03:02:53.663340Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_sample(index):\n    images, targets, image_ids = next(iter(val_data_loader))\n    images = list(img.to(device) for img in images)\n    image = images[index].permute(1,2,0).cpu().numpy()\n    original_image = image.copy()\n\n    boxes = targets[index]['boxes'].numpy()\n    boxes = [np.array(box).astype(np.int32) for box in A.augmentations.bbox_utils.denormalize_bboxes(boxes,image.shape[0],image.shape[1])]\n\n    model.eval()\n    model.to(device)\n    with torch.no_grad():\n        outputs = model(images)\n    pred_classes = outputs['pred_logits'][index].softmax(1).cpu().numpy()\n    pred_boxes = outputs['pred_boxes'][index].cpu().numpy()\n    pred_boxes = [np.array(box).astype(np.int32) for box in A.augmentations.bbox_utils.denormalize_bboxes(pred_boxes,image.shape[0],image.shape[1])]\n\n    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))\n\n    for box in boxes:\n        cv2.rectangle(image,\n                      (box[0], box[1]),\n                      (box[2]+box[0], box[3]+box[1]),\n                      (255, 0, 0), 1)\n    class_color = {0:(255, 255, 0),1:(0, 0, 255),2:(0, 255, 0)}\n    for pred_class, pred_box in zip(pred_classes, pred_boxes):\n        c = pred_class.argmax(0)\n        if c != 3:\n            pred_prob = pred_class.max(0)\n            cv2.rectangle(image,\n                  (pred_box[0], pred_box[1]),\n                  (pred_box[2]+pred_box[0], pred_box[3]+pred_box[1]),\n                  class_color[c], 1)\n            cv2.rectangle(image,\n                         (pred_box[0]+33,pred_box[1]),\n                         (pred_box[0], pred_box[1]-15),\n                          class_color[c], -1)\n            ax2.text(pred_box[0], pred_box[1]-8, f'Class:{c}',color='w',fontsize=9)\n            ax2.text(pred_box[0], pred_box[1], f'Prob:{int(pred_prob*100)}%',color='w',fontsize=9)\n    ax1.set_axis_off()\n    ax1.imshow(original_image)\n    ax2.set_axis_off()\n    ax2.imshow(image)","metadata":{"execution":{"iopub.status.busy":"2022-03-10T03:18:58.757226Z","iopub.execute_input":"2022-03-10T03:18:58.757892Z","iopub.status.idle":"2022-03-10T03:18:58.775778Z","shell.execute_reply.started":"2022-03-10T03:18:58.757854Z","shell.execute_reply":"2022-03-10T03:18:58.775001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = DETRModel(num_classes=num_classes,num_queries=num_queries)\nmodel.load_state_dict(torch.load(\"./detr_best.pth\"))\nshow_sample(0)","metadata":{"execution":{"iopub.status.busy":"2022-03-10T03:19:41.884862Z","iopub.execute_input":"2022-03-10T03:19:41.885132Z","iopub.status.idle":"2022-03-10T03:19:45.542680Z","shell.execute_reply.started":"2022-03-10T03:19:41.885104Z","shell.execute_reply":"2022-03-10T03:19:45.541036Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}