# ---
# 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
# ---

from utils import configurations, utils, base, image_sequence, sentence
import tensorflow as tf
import numpy as np

configs = configurations.DatasetConfigs()
configs.n_features = 128
configs.batch_size = 12
configs.video.features_folder = "rgb_vgg_fc7_features/"
configs.describe()

with configs:
    train_ds, valid_ds, test_ds, embedding_matrix = base.preprocess_datasets()

    video_layer = image_sequence.VideoLayer()
    sentence_layer = sentence.SentenceLayer(embedding_matrix)
    #     audio_layer = get_audio_layer()

    moment = base.MomentVideo(video_layer, sentence_layer, configs.batch_size)

tf.config.set_soft_device_placement(True)

moment.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=5e-4))

moment.evaluate(valid_ds)

with tf.device("GPU:0"):
    moment.fit(train_ds.repeat(), epochs=10, validation_data=valid_ds, steps_per_epoch=600)

from utils.metrics import mIOU, prediction_by_moving_average, intersection_over_union

data = iter(test_ds).get_next()

mean_15 = 0
mean_7 = 0
mean_30 = 0
count = 0
for data in iter(test_ds):
    videos, sentences, y_true = data

    videos_repr = video_layer(videos)
    sentences_repr = sentence_layer(sentences)

    scores_videos = []
    for i in range(configs.batch_size):
        coattention_matrix = matrix_cosine_similarity(videos_repr[i], sentences_repr[i])
        scores_video = get_video_score(videos_repr[i], sentences_repr[i], coattention_matrix)
        scores_videos.append(scores_video)

    scores_videos = tf.stack(scores_videos)
    miou_score = tf.numpy_function(mIOU, [scores_videos, y_true, 7], tf.float64)
    mean_7 += miou_score
    miou_score = tf.numpy_function(mIOU, [scores_videos, y_true, 15], tf.float64)
    mean_15 += miou_score
    miou_score = tf.numpy_function(mIOU, [scores_videos, y_true, 30], tf.float64)
    mean_30 += miou_score
    count += 1

mean_15 = 0
mean_7 = 0
mean_30 = 0
count = 0
for data in iter(train_ds):
    videos, sentences, y_true = data

    videos_repr = moment.video_1(videos)
    sentences_repr = moment.sentence_1(sentences)

    scores_videos = []
    for i in range(moment.batch_size):
        #coattention_matrix = self.matrix_cosine_similarity(videos_repr[i], sentences_repr[i])
        scores_video = moment.get_video_score(videos_repr[i], sentences_repr[i])
        scores_videos.append(scores_video)

    scores_videos = tf.stack(scores_videos)
    miou_score = tf.numpy_function(mIOU, [scores_videos, y_true, 7], tf.float64)
    mean_7 += miou_score
    miou_score = tf.numpy_function(mIOU, [scores_videos, y_true, 15], tf.float64)
    mean_15 += miou_score
    miou_score = tf.numpy_function(mIOU, [scores_videos, y_true, 30], tf.float64)
    mean_30 += miou_score
    count += 1

print(mean_7/count)
print(mean_15/count)
print(mean_30/count)

from itertools import product


def generate_proposals(scores):
    n_proposals = scores.shape[1] // 25
    proposals = list(product(range(n_proposals), range(n_proposals)))
    y_pred = []
    for score in scores:
        max_score = 0
        proposal = [0, 0]
        for proposal in proposals:
            if proposal[1] >= proposal[0]:
                final_score = tf.reduce_mean(score[proposal[0]*25:(proposal[1]+1)*25])
                if final_score > max_score:
                    max_score = final_score
                    max_proposal = proposal
                    
        y_pred.append(max_proposal)
    
    return y_pred


total_scores = 0
count = 0
for data in iter(test_ds):
    videos, sentences, y_true = data

    videos_repr = moment.video_1(videos)
    sentences_repr = moment.sentence_1(sentences)

    scores_videos = []
    for i in range(moment.batch_size):
        coattention_matrix = moment.matrix_cosine_similarity(videos_repr[i], sentences_repr[i])
        scores_video = moment.get_video_score(videos_repr[i], sentences_repr[i], coattention_matrix)
        scores_videos.append(scores_video)

    scores_videos = tf.stack(scores_videos)
    
    y_pred = tf.cast(tf.convert_to_tensor(generate_proposals(scores_videos)), tf.int64)
    scores = np.apply_along_axis(
        intersection_over_union, 1, np.concatenate([y_true, y_pred], axis=-1))
    total_scores += tf.reduce_sum(scores)
    count+=12

print(total_scores/count)

# +
from torch import nn
import torch
from torch.autograd import Variable

def xattn_score_t2i(images, captions, cap_lens, opt):
    """
    Images: (n_image, n_regions, d) matrix of images
    Captions: (n_caption, max_n_word, d) matrix of captions
    CapLens: (n_caption) array of caption lengths
    """
    similarities = []
    n_image = images.size(0)
    n_caption = captions.size(0)
    for i in range(n_caption):
        # Get the i-th text description
        n_word = cap_lens[i]
        cap_i = captions[i, :n_word, :].unsqueeze(0).contiguous()
        # --> (n_image, n_word, d)
        cap_i_expand = cap_i.repeat(n_image, 1, 1)
        """
            word(query): (n_image, n_word, d)
            image(context): (n_image, n_regions, d)
            weiContext: (n_image, n_word, d)
            attn: (n_image, n_region, n_word)
        """
        weiContext, attn = func_attention(cap_i_expand, images, opt, smooth=opt.lambda_softmax)
        print()
        cap_i_expand = cap_i_expand.contiguous()
        weiContext = weiContext.contiguous()
        # (n_image, n_word)
        row_sim = cosine_similarity(cap_i_expand, weiContext, dim=2)
        if opt.agg_func == 'LogSumExp':
            row_sim.mul_(opt.lambda_lse).exp_()
            row_sim = row_sim.sum(dim=1, keepdim=True)
            row_sim = torch.log(row_sim)/opt.lambda_lse
        elif opt.agg_func == 'Max':
            row_sim = row_sim.max(dim=1, keepdim=True)[0]
        elif opt.agg_func == 'Sum':
            row_sim = row_sim.sum(dim=1, keepdim=True)
        elif opt.agg_func == 'Mean':
            row_sim = row_sim.mean(dim=1, keepdim=True)
        else:
            raise ValueError("unknown aggfunc: {}".format(opt.agg_func))
        similarities.append(row_sim)

    # (n_image, n_caption)
    similarities = torch.cat(similarities, 1)
    
    return similarities

class ContrastiveLoss(nn.Module):
    """
    Compute contrastive loss
    """
    def __init__(self, margin=0, max_violation=False):
        super(ContrastiveLoss, self).__init__()
        self.margin = margin
        self.max_violation = max_violation

    def forward(self, scores):
        # compute image-sentence score matrix
#         if self.opt.cross_attn == 't2i':
#             scores = xattn_score_t2i(im, s, s_l, self.opt)
#         elif self.opt.cross_attn == 'i2t':
#             scores = xattn_score_i2t(im, s, s_l, self.opt)
#         else:
#             raise ValueError("unknown first norm type:", opt.raw_feature_norm)
        diagonal = scores.diag().view(scores.size(0), 1)
        d1 = diagonal.expand_as(scores)
        d2 = diagonal.t().expand_as(scores)

        # compare every diagonal score to scores in its column
        # caption retrieval
        cost_s = (1 + scores - d1).clamp(min=0)
        # compare every diagonal score to scores in its row
        # image retrieval
        cost_im = (1 + scores - d2).clamp(min=0)

        # clear diagonals
        mask = torch.eye(scores.size(0)) > .5
        I = Variable(mask)
        cost_s = cost_s.masked_fill_(I, 0)
        cost_im = cost_im.masked_fill_(I, 0)

        # keep the maximum violating negative for each query
        if self.max_violation:
            cost_s = cost_s.max(1)[0]
            cost_im = cost_im.max(0)[0]
        return cost_s.sum() + cost_im.sum()


# +
def get_ranking_loss(scores, margin, top_k=8):
    scores_positives = tf.linalg.diag_part(scores)

    shape_negative = list(scores.shape)
    bool_matrix = tf.ones(shape_negative)
    bool_matrix = tf.linalg.set_diag(bool_matrix, tf.zeros(shape_negative[0])) == 1
    shape_negative[-1] = shape_negative[-1] - 1

    scores_negatives = tf.reshape(scores[bool_matrix], shape_negative)
    top_k_scores_negatives = tf.sort(scores_negatives, axis=1, direction="DESCENDING")[:, 0]

    loss = 1 - scores_positives + top_k_scores_negatives
    loss = tf.where(loss < 0.0, 0.0, loss)

    return tf.reduce_sum(loss)


def margin_based_ranking_loss(scores_videos, scores_sentences, margin=1, top_k=8):
    video_loss = get_ranking_loss(scores_videos, margin, top_k)
    sentence_loss = get_ranking_loss(scores_sentences, margin, top_k)
    
    return video_loss + sentence_loss


# -

scores_videos, scores_sentences = moment.call((data[0], data[1]))

ContrastiveLoss()(torch.from_numpy(scores_videos.numpy()))

get_ranking_loss(scores_videos, 1)

scores = torch.from_numpy(scores_videos.numpy())

data = iter(test_ds).get_next()
videos_repr = moment.video_1(data[0])
sentences_repr = moment.sentence_1(data[1])

images = torch.from_numpy(videos_repr.numpy())
captions = torch.from_numpy(sentences_repr.numpy())
cap_lens = torch.from_numpy(lengths.numpy())


def func_attention(query, context, smooth, eps=1e-8):
    """
    query: (n_context, queryL, d)
    context: (n_context, sourceL, d)
    """
    batch_size_q, queryL = query.size(0), query.size(1)
    batch_size, sourceL = context.size(0), context.size(1)


    # Get attention
    # --> (batch, d, queryL)
    queryT = torch.transpose(query, 1, 2)

    # (batch, sourceL, d)(batch, d, queryL)
    # --> (batch, sourceL, queryL)
    attn = torch.bmm(context, queryT)
#     if opt.raw_feature_norm == "softmax":
#         # --> (batch*sourceL, queryL)
#         attn = attn.view(batch_size*sourceL, queryL)
#         attn = nn.Softmax()(attn)
#         # --> (batch, sourceL, queryL)
#         attn = attn.view(batch_size, sourceL, queryL)
#     elif opt.raw_feature_norm == "l2norm":
#         attn = l2norm(attn, 2)
#     elif opt.raw_feature_norm == "clipped_l2norm": #usado normalmente
#         attn = nn.LeakyReLU(0.1)(attn)
#         attn = l2norm(attn, 2)
#     elif opt.raw_feature_norm == "l1norm":
#         attn = l1norm_d(attn, 2)
#     elif opt.raw_feature_norm == "clipped_l1norm":
#         attn = nn.LeakyReLU(0.1)(attn)
#         attn = l1norm_d(attn, 2)
#     elif opt.raw_feature_norm == "clipped":
#         attn = nn.LeakyReLU(0.1)(attn)
#     elif opt.raw_feature_norm == "no_norm":
#         pass
#     else:
#         raise ValueError("unknown first norm type:", opt.raw_feature_norm)
    # --> (batch, queryL, sourceL)
    attn = torch.transpose(attn, 1, 2).contiguous()
    # --> (batch*queryL, sourceL)
    attn = attn.view(batch_size*queryL, sourceL)
    attn = nn.Softmax()(attn*smooth)
    # --> (batch, queryL, sourceL)
    attn = attn.view(batch_size, queryL, sourceL)
    # --> (batch, sourceL, queryL)
    attnT = torch.transpose(attn, 1, 2).contiguous()

    # --> (batch, d, sourceL)
    contextT = torch.transpose(context, 1, 2)
    # (batch x d x sourceL)(batch x sourceL x queryL)
    # --> (batch, d, queryL)
    weightedContext = torch.bmm(contextT, attnT)
    # --> (batch, queryL, d)
    weightedContext = torch.transpose(weightedContext, 1, 2)

    return weightedContext, attnT

# +
"""
Images: (n_image, n_regions, d) matrix of images
Captions: (n_caption, max_n_word, d) matrix of captions
CapLens: (n_caption) array of caption lengths
"""
similarities = []
n_image = images.size(0)
n_caption = captions.size(0)
for i in range(n_caption):
    # Get the i-th text description
    n_word = cap_lens[i]
    cap_i = captions[i, :n_word, :].unsqueeze(0).contiguous()
    # --> (n_image, n_word, d)
    cap_i_expand = cap_i.repeat(n_image, 1, 1)
    """
        word(query): (n_image, n_word, d)
        image(context): (n_image, n_regions, d)
        weiContext: (n_image, n_word, d)
        attn: (n_image, n_region, n_word)
    """
    weiContext, attn = func_attention(cap_i_expand, images, smooth=1.)
    cap_i_expand = cap_i_expand.contiguous()
    weiContext = weiContext.contiguous()
    # (n_image, n_word)
    row_sim = cosine_similarity_pytorch(cap_i_expand, weiContext, dim=2)
#     print(row_sim)
    row_sim.mul_(1.).exp_()
    row_sim = row_sim.sum(dim=1, keepdim=True)
    row_sim = torch.log(row_sim)/1.
#     elif opt.agg_func == 'Max':
#         row_sim = row_sim.max(dim=1, keepdim=True)[0]
#     elif opt.agg_func == 'Sum':
#         row_sim = row_sim.sum(dim=1, keepdim=True)
#     elif opt.agg_func == 'Mean':
#         row_sim = row_sim.mean(dim=1, keepdim=True)
#     else:
#         raise ValueError("unknown aggfunc: {}".format(opt.agg_func))
    similarities.append(row_sim)

# (n_image, n_caption)
similarities = torch.cat(similarities, 1)
print(similarities)

# +
attn = torch.transpose(attn, 1, 2).contiguous()
# --> (batch*queryL, sourceL)
attn = attn.view(batch_size*queryL, sourceL)
attn = nn.Softmax()(attn*smooth)
# --> (batch, queryL, sourceL)
attn = attn.view(batch_size, queryL, sourceL)
# --> (batch, sourceL, queryL)
attnT = torch.transpose(attn, 1, 2).contiguous()

# --> (batch, d, sourceL)
contextT = torch.transpose(context, 1, 2)
# (batch x d x sourceL)(batch x sourceL x queryL)
# --> (batch, d, queryL)
weightedContext = torch.bmm(contextT, attnT)
# --> (batch, queryL, d)
weightedContext = torch.transpose(weightedContext, 1, 2)


# -

def cosine_similarity(tensor1, tensor2, axis=2):
    num = tf.reduce_sum(tensor1 * tensor2, axis=axis)
    den = tf.norm(tensor1, axis=axis) * tf.norm(tensor2, axis=axis)
    return (num/(den+1e-15))


def cosine_similarity_pytorch(x1, x2, dim=1, eps=1e-8):
    """Returns cosine similarity between x1 and x2, computed along dim."""
    w12 = torch.sum(x1 * x2, dim)
    w1 = torch.norm(x1, 2, dim)
    w2 = torch.norm(x2, 2, dim)
    return (w12 / (w1 * w2).clamp(min=eps)).squeeze()


def attention(query, context, smooth=1., axis=2):
    batch_size_q, queryL = query.shape[0], query.shape[1]
    batch_size, sourceL = context.shape[0], context.shape[1]
    
    queryT = tf.transpose(query, [0, 2, 1])
    attn = tf.matmul(context, queryT)
    attn = tf.transpose(attn, [0, 2, 1])
    attn = tf.reshape(attn, (batch_size*queryL, sourceL))
    attn = tf.nn.softmax(attn*smooth)
    attn = tf.reshape(attn, (batch_size, queryL, sourceL))
    attnT = tf.transpose(attn, [0, 2, 1])
    contextT = tf.transpose(context, [0, 2, 1])
    weighted_context  = tf.matmul(contextT, attnT)
    weighted_context = tf.transpose(weighted_context, [0, 2, 1])
    
    return weighted_context


# +
videos, sentences, lengths, labels = iter(test_ds).get_next()

videos_repr = moment.video_1(videos)
sentences_repr = moment.sentence_1(sentences)

scores = []

for i in range(moment.batch_size):
    n_word = lengths[i]
    cap_i = sentences_repr[i, :n_word, :]
    cap_i_expand = tf.broadcast_to(cap_i, [configs.batch_size, n_word, configs.n_features])
    
    attended_videos = attention(cap_i_expand, videos_repr, axis=2)
    similarity = cosine_similarity(attended_videos, cap_i_expand)
#     print(similarity)
    score = tf.reduce_logsumexp(similarity, axis=1)
    scores.append(score)

scores = tf.stack(scores, axis=1)
print(scores)

# +
videos, sentences, lengths, labels = iter(test_ds).get_next()

videos_repr = moment.video_1(videos)
sentences_repr = moment.sentence_1(sentences)

scores_videos = []
for i in range(moment.batch_size):
    _scores_videos = []
    
    
    for j in range(moment.batch_size):
        coattention_matrix = moment.matrix_cosine_similarity(videos_repr[i], sentences_repr[j])
        scores_video = moment.get_video_score(videos_repr[i], sentences_repr[j], coattention_matrix)

        _scores_videos.append(tf.reduce_logsumexp(scores_video))

    scores_videos.append(_scores_videos)

scores_videos = tf.stack(scores_videos)
# -

scores_videos

StarcoderdataJupyterScriptsDedupFiltered