"""
# 🐠 Reef - Pytorch Starter - FasterRCNN Train

## A self-contained, simple, pure pytorch 🔥 Faster R-CNN implementation with `LB=0.413`

![](https://storage.googleapis.com/kaggle-competitions/kaggle/31703/logos/header.png)

#### FasterR-CNN is one of the SOTA models for Object detection.

### In this notebook we present a simple solution using a pure pytorch Faster R-CNN with pretrained weights, and finetuning it for few epochs.

It is an adapted version of [this notebook](https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train) mentioned in [this comment](https://www.kaggle.com/c/tensorflow-great-barrier-reef/discussion/290016).

## You can find the [inference notebook here](https://www.kaggle.com/julian3833/coral-reef-pytorch-fasterrcnn-infer-0-xxx).

## Details: 
- FasterRCNN from torchvision
- Use Resnet50 backbone

**Update**: Added simple train/validation split in this version, using the "subsequence" split of this notebook: [🐠 Reef - CV strategy: subsequences!](https://www.kaggle.com/julian3833/reef-cv-strategy-subsequences)

Still dropping all the images with no objects, as the model doesn't support them out-of-the-box. The other starters are removing the empty images as well, so it might be a general condition of Object Detection. I'm quite noob in the field to be honest.

# Please, _DO_ upvote if you find this useful!!


&nbsp;
&nbsp;
&nbsp;

#### Changelog

| Version | Description| Dataset| Best LB |
| --- | ----| --- | --- |
| [**V8**](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-train-lb-0-293?scriptVersionId=80517118)  | 2 epochs - Save last epoch | [coral-reef-pytorch-starter-fasterrcnn-weights](https://www.kaggle.com/julian3833/coral-reef-pytorch-starter-fasterrcnn-weights)| `0.293`|
| [**V16**](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-train-lb-0-293?scriptVersionId=80601095) | 4 epochs - Save all epochs | [reef-starter-torch-fasterrcnn-4e](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-4e)| `0.361` |
| [**V17**](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-train-lb-0-369?scriptVersionId=80604402) | Add **validation**. 95-5 split. 8 epochs, keeping track of validation loss. | [reef-starter-torch-fasterrcnn-8e](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-8e)| `0.369` |
| [**V19**](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-train-lb-0-369?scriptVersionId=80610403) | 12 epochs, lower LR | [reef-starter-torch-fasterrcnn-12e](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-12e)| `0.413` |
| [**V24**](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-train-lb-0-369) | V19 with 90-10 train-validation split. Tidy up code. Add Flip. Correct problem with augmentations. | -- | `??` |



"""
"""
# Imports
"""
# Very few imports. This is a pure torch solution!
import cv2
import time

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt

import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2

import torch
import torchvision
from torch.utils.data import DataLoader, Dataset
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection import FasterRCNN
"""
# Constants
"""
DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
BASE_DIR = "../input/tensorflow-great-barrier-reef/train_images/"

NUM_EPOCHS = 12
"""
# Load `df`

### See: [🐠 Reef - CV strategy: subsequences!](https://www.kaggle.com/julian3833/reef-cv-strategy-subsequences)
"""
import ast
df = pd.read_csv("../input/reef-cv-strategy-subsequences-dataframes/train-validation-split/train-0.1.csv")

# string으로 되어있는 annotation을 list of dictionaris로 변환 (사실 train dataset 만들때 미리 해둠)
df['annotations'] = df['annotations'].apply(ast.literal_eval)

# Create the image path for the row
df['image_path'] = "video_" + df['video_id'].astype(str) + "/" + df['video_frame'].astype(str) + ".jpg"

df.head()
# is_Train = True -> df_train / False -> df_val
df_train, df_val = df[df['is_train']], df[~df['is_train']]
df_train[df_train.annotations.str.len()== 0]
# The model doesn't support images with no annotations
# It raises an error that suggest that it just doesn't support them:
# ValueError: No ground-truth boxes available for one of the images during training
# I'm dropping those images for now
# https://discuss.pytorch.org/t/fasterrcnn-images-with-no-objects-present-cause-an-error/117974/3

# annotations가 없으면 모델이 에러나기 때문에 annotations의 길이가 0이상인 것만 남깁니다.
df_train = df_train[df_train.annotations.str.len() > 0 ].reset_index(drop=True)
df_val = df_val[df_val.annotations.str.len() > 0 ].reset_index(drop=True)
df_train.shape[0], df_val.shape[0]
df_train.head()
row = df_train.iloc[0]
boxes = pd.DataFrame(row['annotations'], columns=['x', 'y', 'width', 'height']).astype(float).values
boxes[:, 2] = boxes[:, 0] + boxes[:, 2] # x_max
boxes[:, 3] = boxes[:, 1] + boxes[:, 3] # y_max
print("[x_min, y_min, x_max, y_max]",boxes) 
box_outside_image = ((boxes[:, 0] < 0).any() or (boxes[:, 1] < 0).any() 
                             or (boxes[:, 2] > 1280).any() or (boxes[:, 3] > 720).any())
print(box_outside_image)
"""
# Dataset class
"""
class ReefDataset:

    def __init__(self, df, transforms=None):
        self.df = df
        self.transforms = transforms

    def can_augment(self, boxes):
        """ Check if bounding boxes are OK to augment
        augmentation이 가능한지 확인 하는 함수입니다. annotation이 image의 영역 밖으로 나가지 않도록 합니다.
        
        For example: image_id 1-490 has a bounding box that is partially outside of the image
        It breaks albumentation
        Here we check the margins are within the image to make sure the augmentation can be applied
        """
        
        box_outside_image = ((boxes[:, 0] < 0).any() or (boxes[:, 1] < 0).any() 
                             or (boxes[:, 2] > 1280).any() or (boxes[:, 3] > 720).any())
        return not box_outside_image

    def get_boxes(self, row):
        """
        3D foramt의 bboxes를 반환합니다.
        Returns the bboxes for a given row as a 3D matrix with format [x_min, y_min, x_max, y_max]
        """
        
        boxes = pd.DataFrame(row['annotations'], columns=['x', 'y', 'width', 'height']).astype(float).values
        
        # Change from [x_min, y_min, w, h] to [x_min, y_min, x_max, y_max]
        boxes[:, 2] = boxes[:, 0] + boxes[:, 2] # x_max
        boxes[:, 3] = boxes[:, 1] + boxes[:, 3] # y_max
        return boxes
    
    def get_image(self, row):
        """Gets the image for a given row"""
        
        image = cv2.imread(f'{BASE_DIR}/{row["image_path"]}', cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        image /= 255.0 # normalization
        return image
    
    def __getitem__(self, i):

        row = self.df.iloc[i]
        image = self.get_image(row)
        boxes = self.get_boxes(row)
        
        n_boxes = boxes.shape[0]
        
        # Calculate the area
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]) # (x_max - x_min) * (y_max - y_min)
        
        
        target = {
            'boxes': torch.as_tensor(boxes, dtype=torch.float32),
            'area': torch.as_tensor(area, dtype=torch.float32),
            
            'image_id': torch.tensor([i]),
            
            # There is only one class
            'labels': torch.ones((n_boxes,), dtype=torch.int64),
            
            # Suppose all instances are not crowd
            'iscrowd': torch.zeros((n_boxes,), dtype=torch.int64)            
        }

        if self.transforms and self.can_augment(boxes): #transform이 있고 augmentation이 가능하다면...
            sample = {
                'image': image,
                'bboxes': target['boxes'],
                'labels': target['labels']
            }
            sample = self.transforms(**sample)
            image = sample['image']
            
            if n_boxes > 0: # 타겟이 여러 개일 경우 stack함수를 이용해 쌓아올림 map함수로 tensor 변환
                target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)
        else: # augmentation이 안된다면 그대로 tensor로 변환
            image = ToTensorV2(p=1.0)(image=image)['image']

        return image, target

    def __len__(self):
        return len(self.df)
"""
## Augmentations
"""
# train시에만 augmentation 진행
def get_train_transform():
    return A.Compose([
        A.Flip(0.5),
        ToTensorV2(p=1.0)
    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})


def get_valid_transform():
    return A.Compose([
        ToTensorV2(p=1.0)
    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})
# Define datasets
ds_train = ReefDataset(df_train, get_train_transform())
ds_val = ReefDataset(df_val, get_valid_transform())
"""
## Check one sample
"""
# Let's get an interesting one ;)
df_train[df_train.annotations.str.len() > 12].head()
image, targets = ds_train[2200]
image
targets
boxes = targets['boxes'].cpu().numpy().astype(np.int32)
img = image.permute(1,2,0).cpu().numpy()
fig, ax = plt.subplots(1, 1, figsize=(16, 8))

for box in boxes:
    cv2.rectangle(img,
                  (box[0], box[1]),
                  (box[2], box[3]),
                  (220, 0, 0), 3)
    
ax.set_axis_off()
ax.imshow(img);
"""
## DataLoaders
"""
def collate_fn(batch):
    return tuple(zip(*batch))

dl_train = DataLoader(ds_train, batch_size=8, shuffle=False, num_workers=4, collate_fn=collate_fn)
dl_val = DataLoader(ds_val, batch_size=8, shuffle=False, num_workers=4, collate_fn=collate_fn)
"""
# Create the model
"""
def get_model():
    # load a model; pre-trained on COCO
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

    num_classes = 2  # 1 class (starfish) + background

    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features

    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    model.to(DEVICE)
    return model

model = get_model()
"""
# Train
"""
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.0025, momentum=0.9, weight_decay=0.0005)
# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
lr_scheduler = None

n_batches, n_batches_val = len(dl_train), len(dl_val)
validation_losses = []


for epoch in range(NUM_EPOCHS):
    time_start = time.time()
    loss_accum = 0
    
    for batch_idx, (images, targets) in enumerate(dl_train, 1):
         
        images = list(image.to(DEVICE) for image in images)
        targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]

        # Predict
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        loss_value = losses.item()

        loss_accum += loss_value

        # Back-prop
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

    
    # update the learning rate
    if lr_scheduler is not None:
        lr_scheduler.step()

    # Validation 
    val_loss_accum = 0
        
    # Validation 
    with torch.no_grad():
        for batch_idx, (images, targets) in enumerate(dl_val, 1):
            images = list(image.to(DEVICE) for image in images)
            targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]
            
            val_loss_dict = model(images, targets)
            val_batch_loss = sum(loss for loss in val_loss_dict.values())
            val_loss_accum += val_batch_loss.item()
    
    # Logging
    val_loss = val_loss_accum / n_batches_val
    train_loss = loss_accum / n_batches
    validation_losses.append(val_loss)
    
    # Save model
    chk_name = f'fasterrcnn_resnet50_fpn-e{epoch}.bin'
    torch.save(model.state_dict(), chk_name)
    
    
    elapsed = time.time() - time_start
    
    print(f"[Epoch {epoch+1:2d} / {NUM_EPOCHS:2d}] Train loss: {train_loss:.3f}. Val loss: {val_loss:.3f} --> {chk_name}  [{elapsed:.0f} secs]")   
validation_losses
np.argmin(validation_losses)
"""
# Check result
"""
idx = 0

images, targets = next(iter(dl_val))
images = list(img.to(DEVICE) for img in images)
targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]

boxes = targets[idx]['boxes'].cpu().numpy().astype(np.int32)
sample = images[idx].permute(1,2,0).cpu().numpy()

model.eval()

outputs = model(images)
outputs = [{k: v.detach().cpu().numpy() for k, v in t.items()} for t in outputs]
fig, ax = plt.subplots(1, 1, figsize=(16, 8))

# Red for ground truth
for box in boxes:
    cv2.rectangle(sample,
                  (box[0], box[1]),
                  (box[2], box[3]),
                  (220, 0, 0), 3)

    
# Green for predictions
# Print the first 5
for box in outputs[idx]['boxes'][:5]:
    cv2.rectangle(sample,
                  (box[0], box[1]),
                  (box[2], box[3]),
                  (0, 220, 0), 3)

ax.set_axis_off()
ax.imshow(sample);
"""
# Please, _DO_ upvote if you found it useful!
"""