Upload 5 files

data.py

@@ -0,0 +1,44 @@
+import os
+from PIL import Image
+from torchvision import transforms
+from torch.utils.data import Dataset
+
+def find_mask_file(image_path, mask_dir, mask_extensions=['.png', '.jpg', '.jpeg']):
+    base_name = os.path.splitext(os.path.basename(image_path))[0]
+    for ext in mask_extensions:
+        mask_path = os.path.join(mask_dir, base_name + ext)
+        if os.path.exists(mask_path):
+            return mask_path
+    return None
+
+class SegmentationDataset(Dataset):
+    def __init__(self, image_dir, mask_dir, transform=None):
+        self.image_dir = image_dir
+        self.mask_dir = mask_dir
+        self.transform = transform
+        self.image_filenames = os.listdir(image_dir)
+
+    def __len__(self):
+        return len(self.image_filenames)
+    
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.image_dir, self.image_filenames[idx])
+        mask_path = find_mask_file(img_path, self.mask_dir)
+        image = Image.open(img_path).convert("RGB")
+        mask = Image.open(mask_path).convert("L")
+
+        if self.transform:
+            image = self.transform(image)
+            mask = self.transform(mask)
+
+        return image, mask
+
+def transform_img():
+    transform = transforms.Compose([
+        transforms.Resize((128, 128)),
+        transforms.ToTensor()
+    ])
+    return transform
+
+if __name__ == "__main__":
+    print("Dataset class")

inference.py

@@ -0,0 +1,37 @@
+import torch
+import numpy as np
+from PIL import Image
+from unet import UNet
+from data import transform_img 
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def load_model(weights_path, device):
+    model = UNet(in_channels=3, out_channels=1)  
+    model.load_state_dict(torch.load(weights_path, map_location=device))
+    model.to(device)
+    model.eval()
+    return model
+
+def preprocess_image(image_path):
+    transform = transform_img()
+    image = Image.open(image_path).convert("RGB")
+    return transform(image).unsqueeze(0) 
+
+def predict(model, image_tensor, device):
+    with torch.no_grad():
+        image_tensor = image_tensor.to(device)
+        output = model(image_tensor)
+        output = torch.sigmoid(output) 
+    return output.squeeze(0).cpu().numpy()
+    
+def save_output(mask, save_path):
+    mask = (mask > 0.5).astype(np.uint8)*255 
+    mask_image = Image.fromarray(mask[0])  
+    mask_image.save(save_path)
+
+weights_path = "unet_model.pth"
+model = load_model(weights_path, device)
+image_tensor = preprocess_image("DUTS-TE-Image/ILSVRC2012_test_00000003.jpg")
+mask = predict(model, image_tensor, device)
+save_output(mask, "predicted_mask.jpg")

main.py

@@ -0,0 +1,36 @@
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import torch
+from inference import load_model, preprocess_image, predict 
+
+original_img = Image.open("DUTS-TR-Image/ILSVRC2012_test_00000645.jpg").convert("RGB")
+
+background_with_text = original_img.copy()
+draw = ImageDraw.Draw(background_with_text)
+font_size = 50
+font = ImageFont.truetype("/usr/share/fonts/truetype/freefont/FreeSansBold.ttf", font_size)
+text = "Hello, world!"
+text_position = (50, 50)
+text_color = (255, 255, 255)
+draw.text(text_position, text, fill=text_color, font=font)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+weights_path = "unet_model.pth"
+model = load_model(weights_path, device)
+image_tensor = preprocess_image("DUTS-TR-Image/ILSVRC2012_test_00000645.jpg")
+mask = predict(model, image_tensor, device)
+
+print(mask.shape)
+
+mask = mask.squeeze(0)
+mask_binary = (mask > 0.5).astype(np.uint8) * 255
+mask_img = Image.fromarray(mask_binary, mode="L")
+mask_img = mask_img.resize(original_img.size, resample=Image.NEAREST)
+
+original_rgba = original_img.convert("RGBA")
+
+r, g, b, _ = original_rgba.split()
+subject_img = Image.merge("RGBA", (r, g, b, mask_img))
+
+background_with_text.paste(subject_img, (0, 0), subject_img)
+background_with_text.save("final_output.png")

train.py

@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from unet import UNet
+from torch.utils.data import DataLoader
+from data import SegmentationDataset, transform_img 
+
+transform = transform_img()
+
+train_dataset = SegmentationDataset("DUTS-TR-Image", "DUTS-TR-Mask", transform=transform)
+train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True)
+
+test_dataset = SegmentationDataset("DUTS-TE-Image", "DUTS-TE-Mask", transform=transform)
+test_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = UNet().to(device)
+criterion = nn.BCEWithLogitsLoss()
+optimizer = optim.Adam(model.parameters(), lr=1e-4)
+
+def evaluate_model(model, dataloader, criterion, device):
+    model.eval()
+    total_loss = 0
+    total_correct = 0
+    total_pixels = 0
+
+    with torch.no_grad():
+        for images, masks in dataloader:
+            
+            images = images.to(device)
+            masks = masks.to(device)
+            
+            outputs = model(images)
+            
+            loss = criterion(outputs, masks)
+            total_loss += loss.item()
+
+            preds = torch.sigmoid(outputs) > 0.5
+            total_correct += (preds==masks).sum().item()
+            total_pixels += torch.numel(preds)
+    
+    avg_loss = total_loss / len(dataloader)
+    accuracy = total_correct / total_pixels
+    return avg_loss, accuracy
+
+num_epochs = 2
+total_correct = 0
+total_pixels = 0
+
+train_loss_lst = []
+train_accuracy_lst = []
+test_loss_lst = []
+test_accuracy_lst = []
+
+for epoch in range(num_epochs):
+    print(f"Epoch: {epoch+1}")
+    model.train()
+    epoch_loss = 0
+    
+    for images, masks in train_dataloader:
+        
+        images = images.to(device)
+        masks = masks.to(device)
+        
+        outputs = model(images)
+        
+        loss = criterion(outputs, masks)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        preds = torch.sigmoid(outputs) > 0.5
+        total_correct += (preds==masks).sum().item()
+        total_pixels += torch.numel(preds)
+
+        epoch_loss += loss.item()
+
+    train_accuracy = total_correct / total_pixels        
+    avg_train_loss = epoch_loss/len(train_dataloader)
+    print(f"Train loss at {epoch+1} epoch: {avg_train_loss}")
+    print(f"Train accuracy at {epoch+1} epoch: {train_accuracy}")
+    test_loss, test_accuracy = evaluate_model(model, test_dataloader, criterion, device)
+    print(f"Test loss at {epoch+1} epoch: {test_loss}")
+    print(f"Test accuracy at {epoch+1} epoch: {test_accuracy}")
+    train_loss_lst.append(avg_train_loss)
+    test_loss_lst.append(test_loss)
+    train_accuracy_lst.append(train_accuracy)
+    test_accuracy_lst.append(test_accuracy)

unet.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True)
+        )
+    
+    def forward(self, x):
+        return self.conv(x)
+    
+class UpConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UpConv, self).__init__()
+        self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
+
+    def forward(self, x):
+        return self.up(x)
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super(UNet, self).__init__()
+
+        self.encoder1 = ConvBlock(in_channels, 64)
+        self.encoder2 = ConvBlock(64, 128)
+        self.encoder3 = ConvBlock(128, 256)
+        self.encoder4 = ConvBlock(256, 512)
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        self.bottleneck = ConvBlock(512, 1024)
+
+        self.upconv4 = UpConv(1024, 512)
+        self.decoder4 = ConvBlock(1024, 512)
+        self.upconv3 = UpConv(512, 256)
+        self.decoder3 = ConvBlock(512, 256)
+        self.upconv2 = UpConv(256, 128)
+        self.decoder2 = ConvBlock(256, 128)
+        self.upconv1 = UpConv(128, 64)
+        self.decoder1 = ConvBlock(128, 64)
+
+        self.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        enc1 = self.encoder1(x)
+        enc2 = self.encoder2(self.pool(enc1))
+        enc3 = self.encoder3(self.pool(enc2))
+        enc4 = self.encoder4(self.pool(enc3))
+
+        bottleneck = self.bottleneck(self.pool(enc4))
+
+        dec4 = self.upconv4(bottleneck)
+        dec4 = torch.cat((enc4, dec4), dim=1)
+        dec4 = self.decoder4(dec4)
+        
+        dec3 = self.upconv3(dec4)
+        dec3 = torch.cat((enc3, dec3), dim=1)
+        dec3 = self.decoder3(dec3)
+
+        dec2 = self.upconv2(dec3)
+        dec2 = torch.cat((enc2, dec2), dim=1)
+        dec2 = self.decoder2(dec2)
+
+        dec1 = self.upconv1(dec2)
+        dec1 = torch.cat((enc1, dec1), dim=1)
+        dec1 = self.decoder1(dec1)
+
+        return self.final_conv(dec1)