import os
import gdown
import torch

from networks import U2NET
from utils.saving_utils import save_checkpoint

os.makedirs("prev_checkpoints", exist_ok=True)
gdown.download(
    "https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ",
    "./prev_checkpoints/u2net.pth",
    quiet=False,
)

u_net = U2NET(in_ch=3, out_ch=4)
save_checkpoint(u_net, os.path.join("prev_checkpoints", "u2net_random.pth"))

# u2net.pth contains trained weights
trained_net_pth = os.path.join("prev_checkpoints", "u2net.pth")
# u2net_random.pth contains random weights
custom_net_pth = os.path.join("prev_checkpoints", "u2net_random.pth")

net_state_dict = torch.load(trained_net_pth)
count = 0
for k, v in net_state_dict.items():
    count += 1
print("Total number of layers in trained model are: {}".format(count))

custom_state_dict = torch.load(custom_net_pth)
count = 0
for k, v in custom_state_dict.items():
    count += 1
print("Total number of layers in trained model are: {}".format(count))

total_count = 0
update_count = 0
for k, v in net_state_dict.items():
    total_count += 1
    if custom_state_dict[k].shape == v.shape:
        update_count += 1
        custom_state_dict[k] = v

print(
    "Out of {} layers in custom network, {} layers weights are recovered from trained model".format(
        total_count, update_count
    )
)
torch.save(
    custom_state_dict, os.path.join("prev_checkpoints", "cloth_segm_unet_surgery.pth")
)
print("cloth_segm_unet_surgery.pth is generated in prev_checkpoints directory!")
56a8170b9872223457505feabbe9d4e64ec34aa2
<mask token>

<mask token>
os.makedirs('prev_checkpoints', exist_ok=True)
gdown.download(
    'https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ',
    './prev_checkpoints/u2net.pth', quiet=False)
<mask token>
save_checkpoint(u_net, os.path.join('prev_checkpoints', 'u2net_random.pth'))
<mask token>
for k, v in net_state_dict.items():
    count += 1
print('Total number of layers in trained model are: {}'.format(count))
<mask token>
for k, v in custom_state_dict.items():
    count += 1
print('Total number of layers in trained model are: {}'.format(count))
<mask token>
for k, v in net_state_dict.items():
    total_count += 1
    if custom_state_dict[k].shape == v.shape:
        update_count += 1
        custom_state_dict[k] = v
print(
    'Out of {} layers in custom network, {} layers weights are recovered from trained model'
    .format(total_count, update_count))
torch.save(custom_state_dict, os.path.join('prev_checkpoints',
    'cloth_segm_unet_surgery.pth'))
print('cloth_segm_unet_surgery.pth is generated in prev_checkpoints directory!'
    )

<mask token>
os.makedirs('prev_checkpoints', exist_ok=True)
gdown.download(
    'https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ',
    './prev_checkpoints/u2net.pth', quiet=False)
u_net = U2NET(in_ch=3, out_ch=4)
save_checkpoint(u_net, os.path.join('prev_checkpoints', 'u2net_random.pth'))
trained_net_pth = os.path.join('prev_checkpoints', 'u2net.pth')
custom_net_pth = os.path.join('prev_checkpoints', 'u2net_random.pth')
net_state_dict = torch.load(trained_net_pth)
count = 0
for k, v in net_state_dict.items():
    count += 1
print('Total number of layers in trained model are: {}'.format(count))
custom_state_dict = torch.load(custom_net_pth)
count = 0
for k, v in custom_state_dict.items():
    count += 1
print('Total number of layers in trained model are: {}'.format(count))
total_count = 0
update_count = 0
for k, v in net_state_dict.items():
    total_count += 1
    if custom_state_dict[k].shape == v.shape:
        update_count += 1
        custom_state_dict[k] = v
print(
    'Out of {} layers in custom network, {} layers weights are recovered from trained model'
    .format(total_count, update_count))
torch.save(custom_state_dict, os.path.join('prev_checkpoints',
    'cloth_segm_unet_surgery.pth'))
print('cloth_segm_unet_surgery.pth is generated in prev_checkpoints directory!'
    )

import os
import gdown
import torch
from networks import U2NET
from utils.saving_utils import save_checkpoint
os.makedirs('prev_checkpoints', exist_ok=True)
gdown.download(
    'https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ',
    './prev_checkpoints/u2net.pth', quiet=False)
u_net = U2NET(in_ch=3, out_ch=4)
save_checkpoint(u_net, os.path.join('prev_checkpoints', 'u2net_random.pth'))
trained_net_pth = os.path.join('prev_checkpoints', 'u2net.pth')
custom_net_pth = os.path.join('prev_checkpoints', 'u2net_random.pth')
net_state_dict = torch.load(trained_net_pth)
count = 0
for k, v in net_state_dict.items():
    count += 1
print('Total number of layers in trained model are: {}'.format(count))
custom_state_dict = torch.load(custom_net_pth)
count = 0
for k, v in custom_state_dict.items():
    count += 1
print('Total number of layers in trained model are: {}'.format(count))
total_count = 0
update_count = 0
for k, v in net_state_dict.items():
    total_count += 1
    if custom_state_dict[k].shape == v.shape:
        update_count += 1
        custom_state_dict[k] = v
print(
    'Out of {} layers in custom network, {} layers weights are recovered from trained model'
    .format(total_count, update_count))
torch.save(custom_state_dict, os.path.join('prev_checkpoints',
    'cloth_segm_unet_surgery.pth'))
print('cloth_segm_unet_surgery.pth is generated in prev_checkpoints directory!'
    )

import os
import gdown
import torch

from networks import U2NET
from utils.saving_utils import save_checkpoint

os.makedirs("prev_checkpoints", exist_ok=True)
gdown.download(
    "https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ",
    "./prev_checkpoints/u2net.pth",
    quiet=False,
)

u_net = U2NET(in_ch=3, out_ch=4)
save_checkpoint(u_net, os.path.join("prev_checkpoints", "u2net_random.pth"))

# u2net.pth contains trained weights
trained_net_pth = os.path.join("prev_checkpoints", "u2net.pth")
# u2net_random.pth contains random weights
custom_net_pth = os.path.join("prev_checkpoints", "u2net_random.pth")

net_state_dict = torch.load(trained_net_pth)
count = 0
for k, v in net_state_dict.items():
    count += 1
print("Total number of layers in trained model are: {}".format(count))

custom_state_dict = torch.load(custom_net_pth)
count = 0
for k, v in custom_state_dict.items():
    count += 1
print("Total number of layers in trained model are: {}".format(count))

total_count = 0
update_count = 0
for k, v in net_state_dict.items():
    total_count += 1
    if custom_state_dict[k].shape == v.shape:
        update_count += 1
        custom_state_dict[k] = v

print(
    "Out of {} layers in custom network, {} layers weights are recovered from trained model".format(
        total_count, update_count
    )
)
torch.save(
    custom_state_dict, os.path.join("prev_checkpoints", "cloth_segm_unet_surgery.pth")
)
print("cloth_segm_unet_surgery.pth is generated in prev_checkpoints directory!")