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

# +
import numpy as np
import time
import os
import cv2
import sys

# Now cluster the data

# %pylab inline
# %load_ext autoreload
# %autoreload 2


home_dir = os.path.abspath('../../../../')
if 'eva' in home_dir:
    sys.path.append(home_dir)
else:
    print("Home dir is not what is expected:", home_dir)

import sys; sys.argv=['']; del sys
from loaders.loader_uadetrac import LoaderUADetrac
from eva_storage.external.wnet.WNet import WNet

import torch
import torch.nn as nn
import torch.utils.data
import gc
import config
gc.collect()
torch.cuda.set_device(0)

import eva_storage.external.wnet.wnet_tao.WNet as WNet


# +
import torch.nn.functional as F

vertical_sobel=torch.nn.Parameter(torch.from_numpy(np.array([[[[1,  0,  -1], 
                                            [1,  0,  -1], 
                                            [1,  0,  -1]]]])).float().to(config.device), requires_grad=False)

horizontal_sobel=torch.nn.Parameter(torch.from_numpy(np.array([[[[1,   1,  1], 
                                              [0,   0,  0], 
                                              [-1 ,-1, -1]]]])).float().to(config.device), requires_grad=False)

def gradient_regularization(softmax):
    vert=torch.cat([F.conv2d(softmax[:, i].unsqueeze(1), vertical_sobel) for i in range(softmax.shape[0])], 1)
    hori=torch.cat([F.conv2d(softmax[:, i].unsqueeze(1), horizontal_sobel) for i in range(softmax.shape[0])], 1)
    print('vert', torch.sum(vert))
    print('hori', torch.sum(hori))
    mag=torch.pow(torch.pow(vert, 2)+torch.pow(hori, 2), 0.5)
    mean=torch.mean(mag)
    return mean

def train_op(model, optimizer, input, psi=0.5):
    enc = model(input, returns='enc')
    n_cut_loss=gradient_regularization(enc)*psi
    n_cut_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    dec = model(input, returns='dec')
    rec_loss=torch.mean(torch.pow(torch.pow(input, 2) + torch.pow(dec, 2), 0.5))*(1-psi)
    rec_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return n_cut_loss, rec_loss


# +
batch_size = 10
loader = LoaderUADetrac()
images = loader.load_images(image_size = 256)
X = images.astype(np.float)
train_data = torch.from_numpy(X).float()
num_epochs = 200

N, H, W, C = images.shape
train_data = train_data.permute(0, 3, 1, 2)
assert (train_data.size(1) == 3)
train_loader = torch.utils.data.DataLoader(train_data,
                                           shuffle=True, batch_size=batch_size,
                                           num_workers=1, drop_last=True)

model = WNet.WNet(4)
model = model.to(config.device)
optimizer = torch.optim.SGD(model.parameters(), 0.001)

# -

train_data.shape

# +
for epoch in range(num_epochs):
    st = time.perf_counter()
    for i, images in enumerate(train_loader):
        optimizer.zero_grad()
        images_cuda = images.to(config.device)
        loss1, loss2 = train_op(model, optimizer, images_cuda)

    print('epoch [{}/{}], softcut loss:{:.4f}, reconstruction loss:{:.4f}, time elapsed:{:.4f} (sec)'
              .format(epoch + 1, num_epochs, loss1.data, loss2.data, time.perf_counter() - st))


# -





StarcoderdataJupyterScriptsDedupFiltered