Upload folder using huggingface_hub

f3b050a verified 10 months ago

20.5 kB

	import os
	import numpy as np
	import glob
	import math
	import torch
	import torchvision
	# For everything
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.nn import CrossEntropyLoss, Linear, MSELoss
	from torch.nn import ConvTranspose2d, Conv2d, MaxPool2d, BatchNorm2d
	# For our model
	import torchvision.models as models
	from torchvision import datasets, transforms
	from torchvision.io import read_image
	from torch.utils.data import DataLoader, Dataset
	import torch.optim as optim
	from torch.autograd import Variable
	from torchsummary import summary
	class Nothing(nn.Module):
	def __init__(self):
	super(Nothing,self).__init__()
	def forward(self, radar,satellite):
	return radar, satellite

	class ConvBlock(nn.Module):
	def __init__(self, in_channels, out_channels):
	super(ConvBlock, self).__init__()
	# number of input channels is a number of filters in the previous layer
	# number of output channels is a number of filters in the current layer
	# "same" convolutions
	self.conv = nn.Sequential(
	nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding='same', bias=True),
	nn.BatchNorm2d(out_channels),
	nn.ReLU(inplace=True),
	nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding='same', bias=True),
	nn.BatchNorm2d(out_channels),
	nn.ReLU(inplace=True)
	)

	def forward(self, x):
	x = self.conv(x)
	return x

	class UpConv(nn.Module):
	def __init__(self, in_channels, out_channels):
	super(UpConv, self).__init__()
	self.up = nn.Sequential(
	nn.Upsample(scale_factor=2),
	nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding='same', bias=True),
	nn.BatchNorm2d(out_channels),
	nn.ReLU(inplace=True)
	)
	def forward(self, x):
	x = self.up(x)
	return x
	class AttentionBlock(nn.Module):
	"""Attention block with learnable parameters"""
	def __init__(self, F_g, F_l, n_coefficients):
	"""
	:param F_g: number of feature maps (channels) in previous layer
	:param F_l: number of feature maps in corresponding encoder layer, transferred via skip connection
	:param n_coefficients: number of learnable multi-dimensional attention coefficients
	"""

	super(AttentionBlock, self).__init__()

	self.W_gate = nn.Sequential(
	nn.Conv2d(F_g, n_coefficients, kernel_size=1, stride=1, padding=0, bias=True),
	nn.BatchNorm2d(n_coefficients)
	)

	self.W_x = nn.Sequential(
	nn.Conv2d(F_l, n_coefficients, kernel_size=1, stride=1, padding=0, bias=True),
	nn.BatchNorm2d(n_coefficients)
	)

	self.psi = nn.Sequential(
	nn.Conv2d(n_coefficients, 1, kernel_size=1, stride=1, padding=0, bias=True),
	nn.BatchNorm2d(1),
	nn.Sigmoid()
	)

	self.relu = nn.ReLU(inplace=True)

	def forward(self, gate, skip_connection):
	"""
	:param gate: gating signal from previous layer
	:param skip_connection: activation from corresponding encoder layer
	:return: output activations
	"""
	g1 = self.W_gate(gate)
	x1 = self.W_x(skip_connection)
	psi = self.relu(g1 + x1)
	psi = self.psi(psi)
	out = skip_connection * psi
	return out

	class Recurrent_block(nn.Module):
	def __init__(self,ch_out,t=2):
	super(Recurrent_block,self).__init__()
	self.t = t
	self.ch_out = ch_out
	self.conv = nn.Sequential(
	nn.Conv2d(ch_out,ch_out,kernel_size=3,stride=1,padding='same',bias=True),
	nn.BatchNorm2d(ch_out),
	nn.ReLU(inplace=True)
	)

	def forward(self,x):
	for i in range(self.t):

	if i==0:
	x1 = self.conv(x)

	x1 = self.conv(x+x1)
	return x1

	class RRCNN_block(nn.Module):
	def __init__(self,ch_in,ch_out,t=2):
	super(RRCNN_block,self).__init__()
	self.RCNN = nn.Sequential(
	Recurrent_block(ch_out,t=t),
	Recurrent_block(ch_out,t=t)
	)
	self.Conv_1x1 = nn.Conv2d(ch_in,ch_out,kernel_size=1,stride=1,padding='same')

	def forward(self,x):
	x = self.Conv_1x1(x)
	x1 = self.RCNN(x)
	return x+x1


	class single_conv(nn.Module):
	def __init__(self,ch_in,ch_out):
	super(single_conv,self).__init__()
	self.conv = nn.Sequential(
	nn.Conv2d(ch_in, ch_out, kernel_size=3,stride=1,padding='same',bias=True),
	nn.BatchNorm2d(ch_out),
	nn.ReLU(inplace=True)
	)

	def forward(self,x):
	x = self.conv(x)
	return x
	class Unet(nn.Module):
	def __init__(self, rad_channel=1,sat_channel=1, rad_size=640, sat_size=20):
	super(Unet, self).__init__()
	assert rad_size % sat_size == 0, "rad_size must be divisible by sat_size"
	ratio = rad_size // sat_size
	assert (ratio & (ratio - 1)) == 0, "rad_size/sat_size must be a power of 2"
	self.n_pool = int(math.log2(ratio))
	# Encoder
	self.encoder_blocks = nn.ModuleList()
	self.pools = nn.ModuleList()
	for i in range(self.n_pool):
	in_c = rad_channel * (2**(i))
	out_c = rad_channel * (2**(i+1))
	self.encoder_blocks.append(ConvBlock(in_c, out_c))
	if i < self.n_pool:
	self.pools.append(nn.MaxPool2d(kernel_size=2, stride=2))
	# Bottleneck
	self.mid_conv_1 = single_conv(out_c, out_c)
	self.mid_conv_2 = single_conv(sat_channel, out_c)
	self.mid_merge = ConvBlock(2*out_c, out_c)
	# Decoder
	self.up_convs = nn.ModuleList()
	self.decoder_blocks = nn.ModuleList()
	for i in reversed(range(self.n_pool)):
	up_in = rad_channel * (2**(i+2))
	up_out = rad_channel * (2**(i+1))
	self.up_convs.append(UpConv(up_in, up_out))
	self.decoder_blocks.append(ConvBlock(up_in, up_out))
	self.final_decoder = ConvBlock(4rad_channel, 2rad_channel)
	self.out_conv_R = nn.Conv2d(2*rad_channel, rad_channel, kernel_size=1, padding='same')
	self.out_conv_S = nn.Conv2d(out_c, sat_channel, kernel_size=1, padding='same')
	def forward(self, radar, satellite):
	# Encoding
	enc_feats = []
	x = radar
	for i, block in enumerate(self.encoder_blocks):
	x = block(x)
	enc_feats.append(x)
	if i < self.n_pool:
	x = self.pools[i](x)
	# Bottleneck
	x = F.relu(self.mid_conv_1(x))
	y = F.relu(self.mid_conv_2(satellite))
	x = torch.cat((x, y), dim=1)

	mid_out = self.mid_merge(x)
	pred_sat = self.out_conv_S(mid_out)
	# Decoding
	x = x # input to decoder is original x before mid_merge
	for i in range(self.n_pool):
	x = self.up_convs[i](x)
	x = torch.cat((enc_feats[self.n_pool - 1 - i], x), dim=1)
	x = self.decoder_blocks[i](x)
	x = torch.cat((enc_feats[0], x), dim=1)
	x = self.final_decoder(x)
	pred_rad = self.out_conv_R(x)
	return pred_rad, pred_sat
	# class Unet(nn.Module):
	# def __init__(self,num_channel=1,rad_size=640,sat_size=20):
	# super(Unet, self).__init__()
	# self.MaxPool = nn.MaxPool2d(kernel_size=2, stride=2)
	# self.Conv1 = ConvBlock(1, 2*num_channel)
	# self.Conv2 = ConvBlock(2num_channel, 4num_channel)
	# self.Conv3 = ConvBlock(4num_channel, 8num_channel)
	# self.Conv4 = ConvBlock(8num_channel, 16num_channel)
	# self.Conv5 = ConvBlock(16num_channel, 32num_channel)
	# self.mid_conv_1 = single_conv(32num_channel,32num_channel)
	# self.mid_conv_2 = single_conv(2, 32*num_channel)
	# self.MidConv = ConvBlock(64num_channel, 32num_channel)
	# self.out_conv_S = Conv2d(32*num_channel, 1, (1, 1), padding= 'same')
	# self.Up5 = UpConv(64num_channel, 32num_channel)
	# self.UpConv5 = ConvBlock(64num_channel, 32num_channel)
	# self.Up4 = UpConv(32num_channel, 16num_channel)
	# self.UpConv4 = ConvBlock(32num_channel, 16num_channel)
	# self.Up3 = UpConv(16num_channel, 8num_channel)
	# self.UpConv3 = ConvBlock(16num_channel, 8num_channel)
	# self.Up2 = UpConv(8num_channel, 4num_channel)
	# self.UpConv2 = ConvBlock(8num_channel, 4num_channel)
	# self.Up1 = UpConv(4num_channel, 2num_channel)
	# self.UpConv1 = ConvBlock(4num_channel, 2num_channel)
	# self.out_conv_R = Conv2d(2*num_channel, 1, (1, 1), padding= 'same')
	# def forward(self, radar,satellite):
	# e1 = self.Conv1(radar)
	# e2 = self.MaxPool(e1)
	# e2 = self.Conv2(e2)
	# e3 = self.MaxPool(e2)
	# e3 = self.Conv3(e3)
	# e4 = self.MaxPool(e3)
	# e4 = self.Conv4(e4)
	# e5 = self.MaxPool(e4)
	# e5 = self.Conv5(e5)
	# e6 = self.MaxPool(e5)
	# X = F.relu(self.mid_conv_1(e6))
	# Y = F.relu(self.mid_conv_2(satellite))
	# X = torch.cat((X,Y),1)
	# Y = self.MidConv(X)
	# pred_satellite = self.out_conv_S(Y)
	# d5 = self.Up5(X)
	# d5 = torch.cat((e5, d5), dim=1)
	# d5 = self.UpConv5(d5)
	# d4 = self.Up4(d5)
	# d4 = torch.cat((e4, d4), dim=1)
	# d4 = self.UpConv4(d4)
	# d3 = self.Up3(d4)
	# d3 = torch.cat((e3, d3), dim=1)
	# d3 = self.UpConv3(d3)
	# d2 = self.Up2(d3)
	# d2 = torch.cat((e2, d2), dim=1)
	# d2 = self.UpConv2(d2)
	# d1 = self.Up1(d2)
	# d0 = torch.cat((e1, d1), dim=1)
	# d0 = self.UpConv1(d0)
	# pred_radar = self.out_conv_R(d0)
	# return pred_radar, pred_satellite

	class R2Unet(nn.Module):
	def __init__(self,num_channel=1,t=2):
	super(R2Unet, self).__init__()
	self.MaxPool = nn.MaxPool2d(kernel_size=2, stride=2)
	self.RRCNN1 = RRCNN_block(5,2*num_channel,t=t)
	self.RRCNN2 = RRCNN_block(2num_channel,4num_channel,t=t)
	self.RRCNN3 = RRCNN_block(4num_channel,8num_channel,t=t)
	self.RRCNN4 = RRCNN_block(8num_channel,16num_channel,t=t)
	self.RRCNN5 = RRCNN_block(16num_channel,32num_channel,t=t)
	self.mid_conv_1 = single_conv(32num_channel,32num_channel)
	self.mid_conv_2 = single_conv(2, 32*num_channel)
	self.MidConv = RRCNN_block(64num_channel, 32num_channel)
	self.out_conv_S = Conv2d(32*num_channel, 1, (1, 1), padding= 'same')
	self.Up5 = UpConv(64num_channel, 32num_channel)
	self.UpRRCNN5 = RRCNN_block(64num_channel, 32num_channel)
	self.Up4 = UpConv(32num_channel, 16num_channel)
	self.UpRRCNN4 = RRCNN_block(32num_channel, 16num_channel)
	self.Up3 = UpConv(16num_channel, 8num_channel)
	self.UpRRCNN3 = RRCNN_block(16num_channel, 8num_channel)
	self.Up2 = UpConv(8num_channel, 4num_channel)
	self.UpRRCNN2 = RRCNN_block(8num_channel, 4num_channel)
	self.Up1 = UpConv(4num_channel, 2num_channel)
	self.UpRRCNN1 = RRCNN_block(4num_channel, 2num_channel)
	self.out_conv_R = Conv2d(2*num_channel, 1, (1, 1), padding= 'same')
	def forward(self, radar,satellite):
	e1 = self.RRCNN1(radar)
	e2 = self.MaxPool(e1)
	e2 = self.RRCNN2(e2)
	e3 = self.MaxPool(e2)
	e3 = self.RRCNN3(e3)
	e4 = self.MaxPool(e3)
	e4 = self.RRCNN4(e4)
	e5 = self.MaxPool(e4)
	e5 = self.RRCNN5(e5)
	e6 = self.MaxPool(e5)
	X = F.relu(self.mid_conv_1(e6))
	Y = F.relu(self.mid_conv_2(satellite))
	X = torch.cat((X,Y),1)
	Y = self.MidConv(X)
	pred_satellite = self.out_conv_S(Y)
	d5 = self.Up5(X)
	d5 = torch.cat((e5, d5), dim=1)
	d5 = self.UpRRCNN5(d5)
	d4 = self.Up4(d5)
	d4 = torch.cat((e4, d4), dim=1)
	d4 = self.UpRRCNN4(d4)
	d3 = self.Up3(d4)
	d3 = torch.cat((e3, d3), dim=1)
	d3 = self.UpRRCNN3(d3)
	d2 = self.Up2(d3)
	d2 = torch.cat((e2, d2), dim=1)
	d2 = self.UpRRCNN2(d2)
	d1 = self.Up1(d2)
	d0 = torch.cat((e1, d1), dim=1)
	d0 = self.UpRRCNN1(d0)
	pred_radar = self.out_conv_R(d0)
	return pred_radar, pred_satellite
	class AttUnet(nn.Module):
	def __init__(self,num_channel=1):
	super(AttUnet, self).__init__()
	self.MaxPool = nn.MaxPool2d(kernel_size=2, stride=2)
	self.Conv1 = ConvBlock(5, 2*num_channel)
	self.Conv2 = ConvBlock(2num_channel, 4num_channel)
	self.Conv3 = ConvBlock(4num_channel, 8num_channel)
	self.Conv4 = ConvBlock(8num_channel, 16num_channel)
	self.Conv5 = ConvBlock(16num_channel, 32num_channel)
	self.mid_conv_1 = single_conv(32num_channel,32num_channel)
	self.mid_conv_2 = single_conv(2, 32*num_channel)
	self.MidConv = ConvBlock(64num_channel, 32num_channel)
	self.out_conv_S = Conv2d(32*num_channel, 1, (1, 1), padding= 'same')
	self.Up5 = UpConv(64num_channel, 32num_channel)
	self.Att5 = AttentionBlock(F_g=32num_channel, F_l=32num_channel, n_coefficients=16*num_channel)
	self.UpConv5 = ConvBlock(64num_channel, 32num_channel)
	self.Up4 = UpConv(32num_channel, 16num_channel)
	self.Att4 = AttentionBlock(F_g=16num_channel, F_l=16num_channel, n_coefficients=8*num_channel)
	self.UpConv4 = ConvBlock(32num_channel, 16num_channel)
	self.Up3 = UpConv(16num_channel, 8num_channel)
	self.Att3 = AttentionBlock(F_g=8num_channel, F_l=8num_channel, n_coefficients=4*num_channel)
	self.UpConv3 = ConvBlock(16num_channel, 8num_channel)
	self.Up2 = UpConv(8num_channel, 4num_channel)
	self.Att2 = AttentionBlock(F_g=4num_channel, F_l=4num_channel, n_coefficients=2*num_channel)
	self.UpConv2 = ConvBlock(8num_channel, 4num_channel)
	self.Up1 = UpConv(4num_channel, 2num_channel)
	self.Att1 = AttentionBlock(F_g=2num_channel, F_l=2num_channel, n_coefficients=1*num_channel)
	self.UpConv1 = ConvBlock(4num_channel, 2num_channel)
	self.out_conv_R = Conv2d(2*num_channel, 1, (1, 1), padding= 'same')
	def forward(self, radar,satellite):
	e1 = self.Conv1(radar)
	e2 = self.MaxPool(e1)
	e2 = self.Conv2(e2)
	e3 = self.MaxPool(e2)
	e3 = self.Conv3(e3)
	e4 = self.MaxPool(e3)
	e4 = self.Conv4(e4)
	e5 = self.MaxPool(e4)
	e5 = self.Conv5(e5)
	e6 = self.MaxPool(e5)
	X = F.relu(self.mid_conv_1(e6))
	Y = F.relu(self.mid_conv_2(satellite))
	X = torch.cat((X,Y),1)
	Y = self.MidConv(X)
	pred_satellite = self.out_conv_S(Y)
	d5 = self.Up5(X)
	s4 = self.Att5(gate=d5, skip_connection=e5)
	d5 = torch.cat((s4, d5), dim=1) # concatenate attention-weighted skip connection with previous layer output
	d5 = self.UpConv5(d5)
	d4 = self.Up4(d5)
	s3 = self.Att4(gate=d4, skip_connection=e4)
	d4 = torch.cat((s3, d4), dim=1)
	d4 = self.UpConv4(d4)
	d3 = self.Up3(d4)
	s2 = self.Att3(gate=d3, skip_connection=e3)
	d3 = torch.cat((s2, d3), dim=1)
	d3 = self.UpConv3(d3)
	d2 = self.Up2(d3)
	s1 = self.Att2(gate=d2, skip_connection=e2)
	d2 = torch.cat((s1, d2), dim=1)
	d2 = self.UpConv2(d2)
	d1 = self.Up1(d2)
	s0 = self.Att1(gate=d1, skip_connection=e1)
	d0 = torch.cat((s0, d1), dim=1)
	d0 = self.UpConv1(d0)
	pred_radar = self.out_conv_R(d0)
	return pred_radar, pred_satellite
	class AttR2Unet(nn.Module):
	def __init__(self,num_channel=1,t=2):
	super(AttR2Unet, self).__init__()
	self.MaxPool = nn.MaxPool2d(kernel_size=2, stride=2)
	self.RRCNN1 = RRCNN_block(5, 2*num_channel)
	self.RRCNN2 = RRCNN_block(2num_channel, 4num_channel)
	self.RRCNN3 = RRCNN_block(4num_channel, 8num_channel)
	self.RRCNN4 = RRCNN_block(8num_channel, 16num_channel)
	self.RRCNN5 = RRCNN_block(16num_channel, 32num_channel)
	self.mid_conv_1 = single_conv(32num_channel,32num_channel)
	self.mid_conv_2 = single_conv(2, 32*num_channel)
	self.MidConv = RRCNN_block(64num_channel, 32num_channel)
	self.out_conv_S = Conv2d(32*num_channel, 1, (1, 1), padding= 'same')
	self.Up5 = UpConv(64num_channel, 32num_channel)
	self.Att5 = AttentionBlock(F_g=32num_channel, F_l=32num_channel, n_coefficients=16*num_channel)
	self.UpRRCNN5 = RRCNN_block(64num_channel, 32num_channel)
	self.Up4 = UpConv(32num_channel, 16num_channel)
	self.Att4 = AttentionBlock(F_g=16num_channel, F_l=16num_channel, n_coefficients=8*num_channel)
	self.UpRRCNN4 = RRCNN_block(32num_channel, 16num_channel)
	self.Up3 = UpConv(16num_channel, 8num_channel)
	self.Att3 = AttentionBlock(F_g=8num_channel, F_l=8num_channel, n_coefficients=4*num_channel)
	self.UpRRCNN3 = RRCNN_block(16num_channel, 8num_channel)
	self.Up2 = UpConv(8num_channel, 4num_channel)
	self.Att2 = AttentionBlock(F_g=4num_channel, F_l=4num_channel, n_coefficients=2*num_channel)
	self.UpRRCNN2 = RRCNN_block(8num_channel, 4num_channel)
	self.Up1 = UpConv(4num_channel, 2num_channel)
	self.Att1 = AttentionBlock(F_g=2num_channel, F_l=2num_channel, n_coefficients=1*num_channel)
	self.UpRRCNN1 = RRCNN_block(4num_channel, 2num_channel)
	self.out_conv_R = Conv2d(2*num_channel, 1, (1, 1), padding= 'same')
	def forward(self, radar,satellite):
	e1 = self.RRCNN1(radar)
	e2 = self.MaxPool(e1)
	e2 = self.RRCNN2(e2)
	e3 = self.MaxPool(e2)
	e3 = self.RRCNN3(e3)
	e4 = self.MaxPool(e3)
	e4 = self.RRCNN4(e4)
	e5 = self.MaxPool(e4)
	e5 = self.RRCNN5(e5)
	e6 = self.MaxPool(e5)
	X = F.relu(self.mid_conv_1(e6))
	Y = F.relu(self.mid_conv_2(satellite))
	X = torch.cat((X,Y),1)
	Y = self.MidConv(X)
	pred_satellite = self.out_conv_S(Y)
	d5 = self.Up5(X)
	s4 = self.Att5(gate=d5, skip_connection=e5)
	d5 = torch.cat((s4, d5), dim=1) # concatenate attention-weighted skip connection with previous layer output
	d5 = self.UpRRCNN5(d5)
	d4 = self.Up4(d5)
	s3 = self.Att4(gate=d4, skip_connection=e4)
	d4 = torch.cat((s3, d4), dim=1)
	d4 = self.UpRRCNN4(d4)
	d3 = self.Up3(d4)
	s2 = self.Att3(gate=d3, skip_connection=e3)
	d3 = torch.cat((s2, d3), dim=1)
	d3 = self.UpRRCNN3(d3)
	d2 = self.Up2(d3)
	s1 = self.Att2(gate=d2, skip_connection=e2)
	d2 = torch.cat((s1, d2), dim=1)
	d2 = self.UpRRCNN2(d2)
	d1 = self.Up1(d2)
	s0 = self.Att1(gate=d1, skip_connection=e1)
	d0 = torch.cat((s0, d1), dim=1)
	d0 = self.UpRRCNN1(d0)
	pred_radar = self.out_conv_R(d0)
	return pred_radar, pred_satellite
	class Network(nn.Module):
	def __init__(self,model_type:str,rad_channel:int, sat_channel:int,rad_size:int,sat_size:int):
	super(Network,self).__init__()
	print(model_type)
	if(model_type == "Nothing"):
	self.net = Nothing()
	elif(model_type == "Unet"):
	self.net = Unet(rad_channel=rad_channel,sat_channel=sat_channel,rad_size=rad_size,sat_size=sat_size)
	elif(model_type == "Unet"):
	self.net = Unet(rad_channel=rad_channel,sat_channel=sat_channel,rad_size=rad_size,sat_size=sat_size)
	elif(model_type == "R2Unet"):
	self.net = R2Unet(rad_channel=rad_channel,sat_channel=sat_channel,rad_size=rad_size,sat_size=sat_size)
	elif(model_type == "AttUnet"):
	self.net = AttUnet(rad_channel=rad_channel,sat_channel=sat_channel,rad_size=rad_size,sat_size=sat_size)
	elif(model_type == "AttR2Unet"):
	self.net = AttR2Unet(rad_channel=rad_channel,sat_channel=sat_channel,rad_size=rad_size,sat_size=sat_size)
	else:
	raise ValueError("model_type is wrong")
	def forward(self, radar,satellite):
	pred_radar, pred_satellite = self.net.forward(radar,satellite)
	return pred_radar, pred_satellite