<jupyter_start><jupyter_text>PRIDNet<jupyter_code>class conv_lrelu(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(conv_lrelu, self).__init__()
        self.conv = nn.Sequential(nn.Conv2d(in_ch,out_ch,3, padding = 1),nn.LeakyReLU())

    def forward(self, x):
        x = self.conv(x)
        return x
    
class down(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(down, self).__init__()
        self.conv1 = conv_lrelu(in_ch,out_ch)
        self.conv2 = conv_lrelu(out_ch,out_ch)
        self.down =  nn.MaxPool2d((2,2))
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.down(x)
        return x
    

class up(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(up, self).__init__()       
        self.up =  nn.UpsamplingBilinear2d(scale_factor = 2)
        self.conv1 = conv_lrelu(in_ch,out_ch) 
        self.conv2 = conv_lrelu(out_ch,out_ch) 

    def forward(self, x1, x2):
        x1 = self.up(x1)
        if x1.shape != x2.shape:
            x1 = transforms.functional.resize(x1, x2.shape[2:])
        x = torch.cat([x2, x1], dim=1)
        x = self.conv1(x)
        x = self.conv2(x)
        return x

class UNet(nn.Module):
    def __init__(self, in_ch = 4, CH_PER_SCALE = [32,64,128,256,512], out_ch = 12):
        super(UNet, self).__init__()
        self.inc = conv_lrelu(in_ch, CH_PER_SCALE[0])
        self.inc2 = conv_lrelu(CH_PER_SCALE[0], CH_PER_SCALE[0])
        self.down1 = down(CH_PER_SCALE[0], CH_PER_SCALE[1]) 
        self.down2 = down(CH_PER_SCALE[1],CH_PER_SCALE[2])
        self.down3 = down(CH_PER_SCALE[2],CH_PER_SCALE[3])                
        self.down4 = down(CH_PER_SCALE[3],CH_PER_SCALE[4])                
        self.up1 = up(CH_PER_SCALE[4]+CH_PER_SCALE[3],CH_PER_SCALE[3])
        self.up2 = up(CH_PER_SCALE[3]+CH_PER_SCALE[2],CH_PER_SCALE[2])
        self.up3 = up(CH_PER_SCALE[2]+CH_PER_SCALE[1],CH_PER_SCALE[1])
        self.up4 = up(CH_PER_SCALE[1]+CH_PER_SCALE[0],CH_PER_SCALE[0])
        self.outc = nn.Conv2d(CH_PER_SCALE[0], out_ch, 1, padding = 0)

    def forward(self, x):
        x0 = self.inc(x)
        x0 = self.inc2(x0)
        x1 = self.down1(x0)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x3_up = self.up1(x4,x3)
        x2_up = self.up2(x3_up,x2)
        x1_up = self.up3(x2_up,x1)
        out = self.up4(x1_up,x0)
        out = self.outc(out)
#         out = F.pixel_shuffle(out,2) ## Paper final step rearranges 12 channes to 3 RGB channels
#         out = F.hardtanh(out, min_val=0, max_val=1) #Clamp the top and bottom to 0,1 since pixels can only be in this value
        return out
    
class PRIDNet(nn.Module):
    def __init__(self, in_ch = 4, out_ch = 12):
        super(PRIDNet, self).__init__()
        self.feature_extraction = nn.Sequential(conv_lrelu(in_ch, 32), *[conv_lrelu(32, 32) for i in range(3)])
        self.unet0 = UNet(in_ch = 32, out_ch = 12)
        self.unet1 = UNet(in_ch = 32, out_ch = 12)
        self.unet2 = UNet(in_ch = 32, out_ch = 12)
        self.unet3 = UNet(in_ch = 32, out_ch = 12)
        self.unet4 = UNet(in_ch = 32, out_ch = 12)
        self.avgpool1 = nn.AvgPool2d((2,2))
        self.avgpool2 = nn.AvgPool2d((4,4))
        self.avgpool3 = nn.AvgPool2d((8,8))
        self.avgpool4 = nn.AvgPool2d((16,16))
        self.up4 =  nn.UpsamplingBilinear2d(scale_factor = 16)
        self.up3 =  nn.UpsamplingBilinear2d(scale_factor = 8)
        self.up2 =  nn.UpsamplingBilinear2d(scale_factor = 4)
        self.up1 =  nn.UpsamplingBilinear2d(scale_factor = 2)
        self.out =  nn.Conv2d(32+12*5, out_ch, 1, padding = 0)

    def forward(self, x):
        x_feat = self.feature_extraction(x)
        x0 = self.unet0(x_feat)
        x1 = self.up1(self.unet1(self.avgpool1(x_feat)))
        x2 = self.up2(self.unet2(self.avgpool2(x_feat)))
        x3 = self.up3(self.unet3(self.avgpool3(x_feat)))
        x4 = self.up4(self.unet4(self.avgpool4(x_feat)))
        if x1.shape != x0.shape:
            x1 = transforms.functional.resize(x1, x0.shape[2:])
        if x2.shape != x0.shape:
            x2 = transforms.functional.resize(x2, x0.shape[2:])
        if x3.shape != x0.shape:
            x3 = transforms.functional.resize(x3, x0.shape[2:])
        if x4.shape != x0.shape:
            x4 = transforms.functional.resize(x4, x0.shape[2:])
        x_unet_all = torch.cat([x_feat,x0,x1,x2,x3,x4], axis = 1)
        out = self.out(x_unet_all)
        
        out = F.pixel_shuffle(out,2) ## Paper final step rearranges 12 channes to 3 RGB channels
        out = F.hardtanh(out, min_val=0, max_val=1) #Clamp the top and bottom to 0,1 since pixels can only be in this value
        return out
    
    def load_my_state_dict(self, state_dict):
        own_state = self.state_dict()
        for name, param in state_dict.items():
            if name not in own_state:
                 continue
            #if isinstance(param, self.Parameter):
            else:
                # backwards compatibility for serialized parameters
                param = param.data
            own_state[name].copy_(param)
# model = UNet(in_ch = 4, CH_PER_SCALE = [32,64,128,256,512], out_ch = 12)
# model = model.cuda()
# # model.load_state_dict(torch.load('./results_Sony/sony3830.pth'))
# model.load_state_dict(torch.load('./results_UNet/net_weights/sony3996.pth'))

model = PRIDNet()
model = model.cuda()
model.load_state_dict(torch.load('./results_Sony/sony3830.pth'))
# model.load_state_dict(torch.load('./results_SSIM/net_weights/sony2601.pth'))
with open('./Dataset/Sony_val_raw_list.pickle','rb') as f:
    val_raw_list = pickle.load(f)
val_raw_list = ['./Dataset/'+path for path in val_raw_list]
with open('./Dataset/Sony_val_gt_list.pickle','rb') as f:
    val_gt_list = pickle.load(f)
val_gt_list = ['./Dataset/'+path for path in val_gt_list]

Val_PSNR, Val_SSIM, Val_MSSSIM = validate(model, val_raw_list, val_gt_list,block_size = None, batch_size = 1, 
                                          save_img_dir = './Output_Images_PRIDNet_MSE/')
print(Val_PSNR)
print(Val_SSIM)
print(Val_MSSSIM)<jupyter_output>28.202051
0.7219911
0.8361999
{'fraction_non_alphanumeric': 0.13665803108808292, 'fraction_numerical': 0.04970854922279793, 'mean_word_length': 2.6271285965942455, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '423583', 'license': 'MIT', 'n_tokens_mistral': 2610, 'n_tokens_neox': 2484, 'n_words': 413, 'path': 'Evaluate.ipynb', 'repo_name': 'Jm7752/learn_see_in_dark_comp_vision', 'size': '6176'}