

import torch.nn as nn
import torch
import torchvision
from torchvision import models,transforms
from PIL import Image
from torch.utils.data import DataLoader,Dataset,random_split
import os
import torch.optim as optim
import cv2 as cv
import matplotlib.pyplot as plt
img = cv.imread("./train/images/img_001.png")
img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
plt.imshow(img)
img.shape
class TrainDataset(Dataset):
    def __init__(self):
        self.transform = transforms.Compose([
            transforms.Resize((256,256)),
            transforms.ToTensor()
        ])
    def __len__(self):
        return len(os.listdir("./train/images"))
    def __getitem__(self, index):
        img = Image.open(f"./train/images/img_{index+1:03d}.png")
        mask = Image.open(f"./train/masks/mask_{index+1:03d}.png")
        img,mask = self.transform(img),self.transform(mask)
        mask_binary = (mask>0).float()
        return img, mask_binary
full_ds = TrainDataset()
train_size=int(0.8*len(full_ds))
val_size = len(full_ds) - train_size
train_subset, val_subset=random_split(full_ds,[train_size,val_size])
train_loader = DataLoader(train_subset, batch_size=8, shuffle=True, num_workers=0)
valid_loader = DataLoader(val_subset, batch_size=8, shuffle=False, num_workers=0)
class DoubleConv(nn.Module):
    def __init__(self,in_c, out_c):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_c,out_c,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(out_c),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_c,out_c,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(out_c),
            nn.ReLU(inplace=True),
        )
    def forward(self,x):
        return self.conv(x)

class UNET(nn.Module):
    def __init__(self):
        super().__init__()

        self.enc1 = DoubleConv(3, 64)
        self.pool1 = nn.MaxPool2d(2)

        self.enc2 = DoubleConv(64, 128)
        self.pool2 = nn.MaxPool2d(2)

        self.enc3 = DoubleConv(128, 256)
        self.pool3 = nn.MaxPool2d(2)

        self.enc4 = DoubleConv(256, 512)
        self.pool4 = nn.MaxPool2d(2)

        self.bt = DoubleConv(512, 1024)

        self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.dec4 = DoubleConv(1024, 512)

        self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.dec3 = DoubleConv(512, 256)

        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.dec2 = DoubleConv(256, 128)

        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.dec1 = DoubleConv(128, 64)

        self.out = nn.Conv2d(64, 1, kernel_size=1)

    def forward(self, x):
        enc1 = self.enc1(x)
        p1 = self.pool1(enc1)

        enc2 = self.enc2(p1)
        p2 = self.pool2(enc2)

        enc3 = self.enc3(p2)
        p3 = self.pool3(enc3)

        enc4 = self.enc4(p3)
        p4 = self.pool4(enc4)

        bt = self.bt(p4)

        u4 = self.up4(bt)
        x4 = torch.cat([u4, enc4], dim=1)
        d4 = self.dec4(x4)

        u3 = self.up3(d4)
        x3 = torch.cat([u3, enc3], dim=1)
        d3 = self.dec3(x3)

        u2 = self.up2(d3)
        x2 = torch.cat([u2, enc2], dim=1)
        d2 = self.dec2(x2)

        u1 = self.up1(d2)
        x1 = torch.cat([u1, enc1], dim=1)
        d1 = self.dec1(x1)

        out = self.out(d1)
        return out
model = UNET()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.BCEWithLogitsLoss()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,mode='max')
device = torch.device('cuda')
model = model.to(device)
from torchmetrics import JaccardIndex
iou_metric = JaccardIndex(task="binary").to(device)
decoy = 0
for epoch in range(50):
    model.train()
    total_loss = 0 
    for img,mask in train_loader:
        img = img.to(device)
        mask = mask.to(device)
        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output,mask)
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
    model.eval()
    all_outputs=[]
    all_masks = []
    with torch.no_grad():
        for img, mask in valid_loader:
            img = img.to(device)
            mask = mask.to(device)
            optimizer.zero_grad()
            output = model(img)
            all_outputs.append(output)
            all_masks.append(mask)

    best_iou = 0
    threshold_best = 0.5
    all_outputs = torch.cat(all_outputs,)
    all_masks = torch.cat(all_masks)
    for i in range(10, 65):
        threshold = i / 100
        preds = (torch.sigmoid(all_outputs) > threshold).int()
        iou_metric.reset()
        iou_metric.update(preds, all_masks.int())
        current_iou = iou_metric.compute().item()
        if current_iou > best_iou:
            best_iou = current_iou
            threshold_best = threshold
    scheduler.step(best_iou)
    if best_iou > decoy:
        decoy = best_iou
        torch.save(model.state_dict(), "modelbun.pth")
        best_threshold = threshold_best
    print(f"epoch {epoch+1}, loss {total_loss/len(train_loader)}, iou {best_iou}, cu threshold {threshold_best}")
model = UNET()                     
model.load_state_dict(torch.load("modelbun.pth"))
model = model.to(device)
model.eval()
class TestDataset(Dataset):
    def __init__(self):
        self.transform = transforms.Compose([
            transforms.Resize((256,256)),
            transforms.ToTensor()
        ])
    def __len__(self):
        return len(os.listdir("./test/images"))
    def __getitem__(self, index):
        img = Image.open(f"./test/images/img_{index+1:03d}.png")
        img = self.transform(img)
        return img
test_ds = TestDataset()
test_loader= DataLoader(test_ds, batch_size=8, shuffle=False, num_workers=0)
from torchvision.transforms.functional import to_pil_image
index = 0
model.eval()
with torch.no_grad():
    for img in test_loader:
        img = img.to(device)
        output = model(img)
        for j in range(output.shape[0]):
            # if index == 11 or index == 49:
            #     if index == 11:
            #         mask = (torch.sigmoid(output[j]) > 0.999).float()*255
            #         mask = mask.squeeze(0).cpu().byte()
            #         pil_img = to_pil_image(mask)
            #         pil_img.save(f"C:/Users/David/Desktop/ProblemeML/vanatorul de glitchuri/preds/img_{index+1:03d}.png")
            #         index += 1
            #     else:
            #         mask = (torch.sigmoid(output[j]) > 0.999).float()*255
            #         mask = mask.squeeze(0).cpu().byte()
            #         pil_img = to_pil_image(mask)
            #         pil_img.save(f"C:/Users/David/Desktop/ProblemeML/vanatorul de glitchuri/preds/img_{index+1:03d}.png")
            #         index += 1
            # else:
            #     mask = (torch.sigmoid(output[j]) > best_threshold).float()*255
            #     mask = mask.squeeze(0).cpu().byte()
            #     pil_img = to_pil_image(mask)
            #     pil_img.save(f"C:/Users/David/Desktop/ProblemeML/vanatorul de glitchuri/preds/img_{index+1:03d}.png")
            #     index += 1
            mask = (torch.sigmoid(output[j]) > best_threshold).float()*255
            mask = mask.squeeze(0).cpu().byte()
            pil_img = to_pil_image(mask)
            pil_img.save(f"C:/Users/David/Desktop/ProblemeML/vanatorul de glitchuri/preds/img_{index+1:03d}.png")
            index += 1
