# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %load_ext autoreload
# %autoreload 2
import pdb

# +
import os

import matplotlib
from tqdm import tqdm

from utils.config import opt
from data.dataset import Dataset, TestDataset, inverse_normalize
from model import FasterRCNNVGG16
from torch.autograd import Variable
from torch.utils import data as data_
from trainer import FasterRCNNTrainer
from utils import array_tool as at
from utils.vis_tool import visdom_bbox
from utils.eval_tool import eval_detection_voc

# fix for ulimit
# https://github.com/pytorch/pytorch/issues/973#issuecomment-346405667
import resource

rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (20480, rlimit[1]))

matplotlib.use('agg')
import numpy as np
import cupy as cp

from model.utils.bbox_tools import bbox2loc, bbox_iou, loc2bbox
from model.utils.nms import non_maximum_suppression

from collections import namedtuple
import time
from torch.nn import functional as F
from model.utils.creator_tool import AnchorTargetCreator, ProposalTargetCreator

from torch import nn
import torch as t
from torch.autograd import Variable
from utils import array_tool as at
from utils.vis_tool import Visualizer

from utils.config import opt
# from torchnet.meter import ConfusionMeter, AverageValueMeter

# +
dataset = Dataset(opt)

print('load data')
dataloader = data_.DataLoader(dataset, \
                                  batch_size=1, \
                                  shuffle=True, \
                                  pin_memory=True,\
                                  num_workers=opt.num_workers)
testset = TestDataset(opt)
test_dataloader = data_.DataLoader(testset,
                                       batch_size=1,
                                       num_workers=opt.test_num_workers,
                                       shuffle=False, \
                                       pin_memory=True
                                       )
# -

faster_rcnn = FasterRCNNVGG16()
print('model construct completed')

trainer = FasterRCNNTrainer(faster_rcnn).cuda()

best_map = 0
lr_ = opt.lr

img, bbox_, label_, scale = next(iter(dataloader))

img.shape,bbox_.shape, label_.shape, scale.shape

img, bbox, label = img.cuda().float(), bbox_.cuda(), label_.cuda()

# + deletable=false editable=false run_control={"frozen": true}
# trainer(img, bbox_, label_, scale)
# -

 imgs, bboxes, labels, scale = img, bbox_, label_, scale

# +
n = bboxes.shape[0]
if n != 1:
    raise ValueError('Currently only batch size 1 is supported.')

_, _, H, W = imgs.shape
img_size = (H, W)
# -

features = trainer.faster_rcnn.extractor(imgs)

features.shape

trainer.faster_rcnn.rpn

rpn_locs, rpn_scores, rois, roi_indices, anchor = trainer.faster_rcnn.rpn(features, img_size, scale)

rpn_locs.shape, rpn_scores.shape, rois.shape, roi_indices.shape, anchor.shape

bbox = bboxes[0]
label = labels[0]
rpn_score = rpn_scores[0]
rpn_loc = rpn_locs[0]
roi = rois

sample_roi, gt_roi_loc, gt_roi_label = trainer.proposal_target_creator(
            roi,
            at.tonumpy(bbox), # at = array_tools,tensor to numpy 用不着了，在pytorch0.4里
            at.tonumpy(label),
            trainer.loc_normalize_mean,
            trainer.loc_normalize_std)

sample_roi.shape,gt_roi_loc.shape, gt_roi_label.shape

sample_roi_index = t.zeros(len(sample_roi))

roi_cls_loc, roi_score = trainer.faster_rcnn.head(
    features,
    sample_roi,
    sample_roi_index)

roi_cls_loc.shape, roi_score.shape

gt_rpn_loc, gt_rpn_label = trainer.anchor_target_creator(
            at.tonumpy(bbox),
            anchor,
            img_size)

gt_rpn_loc.shape, gt_rpn_label.shape

from trainer import _fast_rcnn_loc_loss

gt_rpn_label = at.tovariable(gt_rpn_label).long()
gt_rpn_loc = at.tovariable(gt_rpn_loc)

# + run_control={"marked": true}
rpn_loc_loss = _fast_rcnn_loc_loss(rpn_loc, gt_rpn_loc, gt_rpn_label.data,
                                   trainer.rpn_sigma)

# + run_control={"marked": true}
rpn_loc_loss
# -

rpn_score.shape,gt_rpn_label.shape

rpn_cls_loss = F.cross_entropy(rpn_score, gt_rpn_label.cuda(), ignore_index=-1)

_gt_rpn_label = gt_rpn_label[gt_rpn_label > -1]
_rpn_score = at.tonumpy(rpn_score)[at.tonumpy(gt_rpn_label) > -1]
# self.rpn_cm.add(at.totensor(_rpn_score, False), _gt_rpn_label.data.long())

_gt_rpn_label.shape,_rpn_score.shape

# ------------------ ROI losses (fast rcnn loss) -------------------#
n_sample = roi_cls_loc.shape[0]

n_sample

roi_cls_loc.shape

# + run_control={"marked": false}
roi_cls_loc = roi_cls_loc.view(n_sample, -1, 4)
# -

roi_cls_loc.shape

gt_roi_label

# + run_control={"marked": false}
roi_loc = roi_cls_loc[t.arange(0, n_sample).long().cuda(), \
      at.totensor(gt_roi_label).long()]
# -

roi_loc.shape

# + run_control={"marked": false}
gt_roi_label = at.tovariable(gt_roi_label).long()
gt_roi_loc = at.tovariable(gt_roi_loc)
# -



# +
roi_loc_loss = _fast_rcnn_loc_loss(
    roi_loc.contiguous(),
    gt_roi_loc,
    gt_roi_label.data,
    self.roi_sigma)

roi_cls_loss = nn.CrossEntropyLoss()(roi_score, gt_roi_label.cuda())

# self.roi_cm.add(at.totensor(roi_score, False), gt_roi_label.data.long())

losses = [rpn_loc_loss, rpn_cls_loss, roi_loc_loss, roi_cls_loss]
losses = losses + [sum(losses)]

StarcoderdataJupyterScriptsDedupFiltered