Upload 2 files

amodal_completion_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4933917dbfac8f3970d150b7fce00c95c58da724e42f21d74964ea53c13c625a
+size 124272930

final1_2.py

@@ -0,0 +1,981 @@
+# -*- coding: utf-8 -*-
+"""final1.2.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1v6-6x7lqt6gr9VIauNVHIwjvIkewk8eT
+"""
+
+
+
+"""## FINAL 1.2"""
+
+
+
+pip install torchmetrics lpips
+
+# PyTorch, Torchvision
+import torch
+from torch import nn
+from torchvision.transforms import ToPILImage, ToTensor
+from torchvision.utils import make_grid
+from torchvision.io import write_video
+
+# Common
+from pathlib import Path
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+import random
+import json
+from IPython.display import Video
+
+# Utils from Torchvision
+tensor_to_image = ToPILImage()
+image_to_tensor = ToTensor()
+
+def get_img_dict(img_dir):
+    img_files = [x for x in img_dir.iterdir() if x.name.endswith('.png') or x.name.endswith('.tiff')]
+    img_files.sort()
+
+    img_dict = {}
+    for img_file in img_files:
+        img_type = img_file.name.split('_')[0]
+        if img_type not in img_dict:
+            img_dict[img_type] = []
+        img_dict[img_type].append(img_file)
+    return img_dict
+
+
+def get_sample_dict(sample_dir):
+
+    camera_dirs = [x for x in sample_dir.iterdir() if 'camera' in x.name]
+    camera_dirs.sort()
+
+    sample_dict = {}
+
+    for cam_dir in camera_dirs:
+        cam_dict = {}
+        cam_dict['scene'] = get_img_dict(cam_dir)
+
+        obj_dirs = [x for x in cam_dir.iterdir() if 'obj_' in x.name]
+        obj_dirs.sort()
+
+        for obj_dir in obj_dirs:
+            cam_dict[obj_dir.name] = get_img_dict(obj_dir)
+
+        sample_dict[cam_dir.name] = cam_dict
+
+    return sample_dict
+
+!wget https://huggingface.co/datasets/Amar-S/MOVi-MC-AC/resolve/main/test_obj_descriptors.json
+#Download Descriptors, Readme, etc.
+!wget https://huggingface.co/datasets/Amar-S/MOVi-MC-AC/resolve/main/train_obj_descriptors.json
+!wget https://huggingface.co/datasets/Amar-S/MOVi-MC-AC/resolve/main/ex_vis.mp4
+!wget https://huggingface.co/datasets/Amar-S/MOVi-MC-AC/resolve/main/README.md
+!wget "https://huggingface.co/datasets/Amar-S/MOVi-MC-AC/resolve/main/Notice%201%20-%20Unlimited_datasets.pdf"
+!wget https://huggingface.co/datasets/Amar-S/MOVi-MC-AC/resolve/main/.gitattributes
+#Test to see if you are on the right huggingface repo
+from huggingface_hub import HfApi, hf_hub_download
+import random, os
+api = HfApi()
+repo_id = "Amar-S/MOVi-MC-AC"
+# # List all files in the repo
+files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
+# # Separate train and test files
+train_files = [f for f in files if f.startswith("train/") and not f.endswith(".json")]
+test_files = [f for f in files if f.startswith("test/") and not f.endswith(".json")]
+print(f"Found {len(train_files)} train files and {len(test_files)} test files.")
+#Download 4% of Train/Test files
+import os
+import random
+import shutil
+from huggingface_hub import hf_hub_download
+os.makedirs("/content/data/train", exist_ok=True)
+os.makedirs("/content/data/test", exist_ok=True)
+# # Sample 4% of each split (as you were doing)
+subset_train = random.sample(train_files, int(len(train_files) * 0.005))
+subset_test = random.sample(test_files, int(len(test_files) * 0.005))
+# # Download the training files (uncomment and fix)
+for file in subset_train:
+    out_path = hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=file)
+    dest_path = f"/content/data/train/{os.path.basename(file)}"
+    shutil.copyfile(out_path, dest_path)  # COPY the actual file content instead of renaming symlink
+# # Download the test files
+for file in subset_test:
+    out_path = hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=file)
+    dest_path = f"/content/data/test/{os.path.basename(file)}"
+    shutil.copyfile(out_path, dest_path)  # COPY the actual file content here as well
+
+import os
+
+# Untar all files in data/train
+train_dir = "data/train"
+for file in os.listdir(train_dir):
+    if file.endswith(".tar.gz"):
+        filepath = os.path.join(train_dir, file)
+        !tar -xzf {filepath} -C {train_dir}
+
+# Untar all files in data/test
+test_dir = "data/test"
+for file in os.listdir(test_dir):
+    if file.endswith(".tar.gz"):
+        filepath = os.path.join(test_dir, file)
+        !tar -xzf {filepath} -C {test_dir}
+
+
+
+import os
+from pathlib import Path
+root = Path('/content/data')  # or wherever your files live
+deleted = 0
+for archive in root.rglob('*.tar.gz'):
+    try:
+        archive.unlink()
+        print(f"Deleted {archive}")
+        deleted += 1
+    except Exception as e:
+        print(f"Error deleting {archive}: {e}")
+print(f"Total deleted: {deleted}")
+
+pip install torchmetrics lpips
+
+import matplotlib.pyplot as plt
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+import lpips
+import matplotlib.pyplot as plt
+import torch
+
+def visualize_results(model, dataloader, device, num_samples=8):
+    """Visualize results with properly masked output (no background)"""
+    model.eval()
+    samples_shown = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            if samples_shown >= num_samples:
+                break
+
+            rgb = batch['rgb'].to(device)
+            modal_mask = batch['modal_mask'].to(device)
+            amodal_mask = batch['amodal_mask'].to(device)
+            gt_amodal_rgb = batch['amodal_rgb'].to(device)
+
+            input_tensor = torch.cat([rgb, modal_mask, amodal_mask], dim=1)
+            pred = model(input_tensor)
+
+            pred_masked = pred * amodal_mask  # Remove background from prediction
+            gt_masked = gt_amodal_rgb * amodal_mask  # Ensure GT is also masked consistently
+
+            for i in range(rgb.shape[0]):
+                if samples_shown >= num_samples:
+                    break
+
+                fig, axes = plt.subplots(1, 6, figsize=(24, 4))
+
+                # Scene RGB
+                axes[0].imshow(rgb[i].cpu().permute(1, 2, 0))
+                axes[0].set_title('Scene RGB')
+                axes[0].axis('off')
+
+                # Amodal Mask
+                axes[1].imshow(amodal_mask[i, 0].cpu(), cmap='gray')
+                axes[1].set_title('Amodal Mask')
+                axes[1].axis('off')
+
+                # Modal Mask
+                axes[2].imshow(modal_mask[i, 0].cpu(), cmap='gray')
+                axes[2].set_title('Modal Mask')
+                axes[2].axis('off')
+
+                # Ground Truth Amodal RGB (masked)
+                axes[3].imshow(gt_masked[i].cpu().permute(1, 2, 0))
+                axes[3].set_title('GT Amodal RGB')
+                axes[3].axis('off')
+
+                # Predicted Amodal RGB (masked)
+                axes[4].imshow(pred_masked[i].cpu().permute(1, 2, 0))
+                axes[4].set_title('Predicted Amodal RGB')
+                axes[4].axis('off')
+
+                # Difference Heatmap
+                diff = torch.abs(pred_masked[i] - gt_masked[i]).mean(dim=0)
+                im = axes[5].imshow(diff.cpu(), cmap='hot')
+                axes[5].set_title('Prediction Error')
+                axes[5].axis('off')
+                plt.colorbar(im, ax=axes[5])
+
+                plt.tight_layout()
+                plt.show()
+
+                samples_shown += 1
+
+
+
+# STEP 4: Add this function for better evaluation:
+def evaluate_metrics(model, dataloader, device):
+    """Compute evaluation metrics only within object regions"""
+    model.eval()
+    total_mse = 0
+    occluded_mse = 0
+    visible_mse = 0
+    total_pixels = 0
+    occluded_pixels = 0
+    visible_pixels = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            rgb = batch['rgb'].to(device)
+            modal_mask = batch['modal_mask'].to(device)
+            amodal_mask = batch['amodal_mask'].to(device)
+            occluded_mask = batch['occluded_mask'].to(device)
+            gt_amodal_rgb = batch['amodal_rgb'].to(device)
+
+            input_tensor = torch.cat([rgb, modal_mask, amodal_mask], dim=1)
+            pred = model(input_tensor)
+
+            # Mask both prediction and ground truth to object regions only
+            pred_masked = pred * amodal_mask
+            gt_masked = gt_amodal_rgb * amodal_mask
+
+            # Overall MSE within object region
+            object_pixels = amodal_mask.sum()
+            if object_pixels > 0:
+                mse = F.mse_loss(pred_masked, gt_masked, reduction='sum')
+                total_mse += mse.item()
+                total_pixels += object_pixels.item()
+
+            # Occluded region MSE
+            occluded_region = occluded_mask * amodal_mask
+            occ_pixels = occluded_region.sum()
+            if occ_pixels > 0:
+                occ_mse = F.mse_loss(pred_masked * occluded_region,
+                                   gt_masked * occluded_region, reduction='sum')
+                occluded_mse += occ_mse.item()
+                occluded_pixels += occ_pixels.item()
+
+            # Visible region MSE
+            visible_region = modal_mask * amodal_mask
+            vis_pixels = visible_region.sum()
+            if vis_pixels > 0:
+                vis_mse = F.mse_loss(pred_masked * visible_region,
+                                   gt_masked * visible_region, reduction='sum')
+                visible_mse += vis_mse.item()
+                visible_pixels += vis_pixels.item()
+
+    return {
+        'total_mse': total_mse / total_pixels if total_pixels > 0 else 0,
+        'occluded_mse': occluded_mse / occluded_pixels if occluded_pixels > 0 else 0,
+        'visible_mse': visible_mse / visible_pixels if visible_pixels > 0 else 0,
+    }
+
+
+
+def calculate_metrics(model, dataloader, device):
+    """Computes PSNR, SSIM, LPIPS, and IoU between predictions and GT amodal RGBs."""
+
+    model.eval()
+    psnr = PeakSignalNoiseRatio().to(device)
+    ssim = StructuralSimilarityIndexMeasure().to(device)
+    lpips_loss = lpips.LPIPS(net='alex').to(device)
+
+    total_psnr, total_ssim, total_lpips = 0, 0, 0
+    total_iou = 0
+    count = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            rgb = batch['rgb'].to(device)
+            modal_mask = batch['modal_mask'].to(device)
+            amodal_mask = batch['amodal_mask'].to(device)
+            gt_amodal_rgb = batch['amodal_rgb'].to(device)
+
+            input_tensor = torch.cat([rgb, modal_mask, amodal_mask], dim=1)
+            pred = model(input_tensor)
+
+            pred_masked = pred * amodal_mask
+            gt_masked = gt_amodal_rgb * amodal_mask
+
+            for i in range(pred.shape[0]):
+                pred_i = pred_masked[i].unsqueeze(0)
+                gt_i = gt_masked[i].unsqueeze(0)
+
+                # Resize for LPIPS if necessary (it requires >= 64x64)
+                if pred_i.shape[-1] < 64 or pred_i.shape[-2] < 64:
+                    continue
+
+                total_psnr += psnr(pred_i, gt_i).item()
+                total_ssim += ssim(pred_i, gt_i).item()
+                total_lpips += lpips_loss(pred_i, gt_i).item()
+
+                # mIoU between masks
+                intersection = (amodal_mask[i] * (pred[i] > 0.5)).sum()
+                union = ((amodal_mask[i] + (pred[i] > 0.5)) > 0).sum()
+                if union > 0:
+                    iou = intersection.float() / union.float()
+                    total_iou += iou.item()
+
+                count += 1
+
+    if count == 0:
+        return {"psnr": 0, "ssim": 0, "lpips": 0, "miou": 0}
+
+    return {
+        "psnr": total_psnr / count,
+        "ssim": total_ssim / count,
+        "lpips": total_lpips / count,
+        "miou": total_iou / count
+    }
+
+pip install torchmetrics lpips
+
+import matplotlib.pyplot as plt
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+import lpips
+import matplotlib.pyplot as plt
+import torch
+
+def visualize_results(model, dataloader, device, num_samples=8):
+    """Visualize results with properly masked output (no background)"""
+    model.eval()
+    samples_shown = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            if samples_shown >= num_samples:
+                break
+
+            rgb = batch['rgb'].to(device)
+            modal_mask = batch['modal_mask'].to(device)
+            amodal_mask = batch['amodal_mask'].to(device)
+            gt_amodal_rgb = batch['amodal_rgb'].to(device)
+
+            input_tensor = torch.cat([rgb, modal_mask, amodal_mask], dim=1)
+            pred = model(input_tensor)
+
+            pred_masked = pred * amodal_mask  # Remove background from prediction
+            gt_masked = gt_amodal_rgb * amodal_mask  # Ensure GT is also masked consistently
+
+            for i in range(rgb.shape[0]):
+                if samples_shown >= num_samples:
+                    break
+
+                fig, axes = plt.subplots(1, 6, figsize=(24, 4))
+
+                # Scene RGB
+                axes[0].imshow(rgb[i].cpu().permute(1, 2, 0))
+                axes[0].set_title('Scene RGB')
+                axes[0].axis('off')
+
+                # Amodal Mask
+                axes[1].imshow(amodal_mask[i, 0].cpu(), cmap='gray')
+                axes[1].set_title('Amodal Mask')
+                axes[1].axis('off')
+
+                # Modal Mask
+                axes[2].imshow(modal_mask[i, 0].cpu(), cmap='gray')
+                axes[2].set_title('Modal Mask')
+                axes[2].axis('off')
+
+                # Ground Truth Amodal RGB (masked)
+                axes[3].imshow(gt_masked[i].cpu().permute(1, 2, 0))
+                axes[3].set_title('GT Amodal RGB')
+                axes[3].axis('off')
+
+                # Predicted Amodal RGB (masked)
+                axes[4].imshow(pred_masked[i].cpu().permute(1, 2, 0))
+                axes[4].set_title('Predicted Amodal RGB')
+                axes[4].axis('off')
+
+                # Difference Heatmap
+                diff = torch.abs(pred_masked[i] - gt_masked[i]).mean(dim=0)
+                im = axes[5].imshow(diff.cpu(), cmap='hot')
+                axes[5].set_title('Prediction Error')
+                axes[5].axis('off')
+                plt.colorbar(im, ax=axes[5])
+
+                plt.tight_layout()
+                plt.show()
+
+                samples_shown += 1
+
+
+def evaluate_metrics(model, dataloader, device):
+    """Compute evaluation metrics only within object regions"""
+    model.eval()
+    total_mse = 0
+    occluded_mse = 0
+    visible_mse = 0
+    total_pixels = 0
+    occluded_pixels = 0
+    visible_pixels = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            rgb = batch['rgb'].to(device)
+            modal_mask = batch['modal_mask'].to(device)
+            amodal_mask = batch['amodal_mask'].to(device)
+            occluded_mask = batch['occluded_mask'].to(device)
+            gt_amodal_rgb = batch['amodal_rgb'].to(device)
+
+            input_tensor = torch.cat([rgb, modal_mask, amodal_mask], dim=1)
+            pred = model(input_tensor)
+
+            # Mask both prediction and ground truth to object regions only
+            pred_masked = pred * amodal_mask
+            gt_masked = gt_amodal_rgb * amodal_mask
+
+            # Overall MSE within object region
+            object_pixels = amodal_mask.sum()
+            if object_pixels > 0:
+                mse = F.mse_loss(pred_masked, gt_masked, reduction='sum')
+                total_mse += mse.item()
+                total_pixels += object_pixels.item()
+
+            # Occluded region MSE
+            occluded_region = occluded_mask * amodal_mask
+            occ_pixels = occluded_region.sum()
+            if occ_pixels > 0:
+                occ_mse = F.mse_loss(pred_masked * occluded_region,
+                                   gt_masked * occluded_region, reduction='sum')
+                occluded_mse += occ_mse.item()
+                occluded_pixels += occ_pixels.item()
+
+            # Visible region MSE
+            visible_region = modal_mask * amodal_mask
+            vis_pixels = visible_region.sum()
+            if vis_pixels > 0:
+                vis_mse = F.mse_loss(pred_masked * visible_region,
+                                   gt_masked * visible_region, reduction='sum')
+                visible_mse += vis_mse.item()
+                visible_pixels += vis_pixels.item()
+
+    return {
+        'total_mse': total_mse / total_pixels if total_pixels > 0 else 0,
+        'occluded_mse': occluded_mse / occluded_pixels if occluded_pixels > 0 else 0,
+        'visible_mse': visible_mse / visible_pixels if visible_pixels > 0 else 0,
+    }
+
+
+
+def calculate_metrics(model, dataloader, device):
+    """Computes PSNR, SSIM, LPIPS, and IoU between predictions and GT amodal RGBs."""
+
+    model.eval()
+    psnr = PeakSignalNoiseRatio().to(device)
+    ssim = StructuralSimilarityIndexMeasure().to(device)
+    lpips_loss = lpips.LPIPS(net='alex').to(device)
+
+    total_psnr, total_ssim, total_lpips = 0, 0, 0
+    total_iou = 0
+    count = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            rgb = batch['rgb'].to(device)
+            modal_mask = batch['modal_mask'].to(device)
+            amodal_mask = batch['amodal_mask'].to(device)
+            gt_amodal_rgb = batch['amodal_rgb'].to(device)
+
+            input_tensor = torch.cat([rgb, modal_mask, amodal_mask], dim=1)
+            pred = model(input_tensor)
+
+            pred_masked = pred * amodal_mask
+            gt_masked = gt_amodal_rgb * amodal_mask
+
+            for i in range(pred.shape[0]):
+                pred_i = pred_masked[i].unsqueeze(0)
+                gt_i = gt_masked[i].unsqueeze(0)
+
+                # Resize for LPIPS if necessary (it requires >= 64x64)
+                if pred_i.shape[-1] < 64 or pred_i.shape[-2] < 64:
+                    continue
+
+                total_psnr += psnr(pred_i, gt_i).item()
+                total_ssim += ssim(pred_i, gt_i).item()
+                total_lpips += lpips_loss(pred_i, gt_i).item()
+
+                # mIoU between masks
+                intersection = (amodal_mask[i] * (pred[i] > 0.5)).sum()
+                union = ((amodal_mask[i] + (pred[i] > 0.5)) > 0).sum()
+                if union > 0:
+                    iou = intersection.float() / union.float()
+                    total_iou += iou.item()
+
+                count += 1
+
+    if count == 0:
+        return {"psnr": 0, "ssim": 0, "lpips": 0, "miou": 0}
+
+    return {
+        "psnr": total_psnr / count,
+        "ssim": total_ssim / count,
+        "lpips": total_lpips / count,
+        "miou": total_iou / count
+    }
+
+
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from pathlib import Path
+from PIL import Image, ImageChops
+import numpy as np
+
+class ModalAmodalDataset(Dataset):
+    def __init__(self, root_dir, split, img_size=(128, 128), max_samples=None, val_split=0.2, use_val_from_train=False):
+        self.root_dir = Path(root_dir)
+        self.img_size = img_size
+        self.max_samples = max_samples
+        self.val_split = val_split
+        self.use_val_from_train = use_val_from_train
+        self.split = split
+
+        if split == 'val' and use_val_from_train:
+            # Load from train folder but use validation subset
+            self.root_dir = self.root_dir / 'train'
+        else:
+            self.root_dir = self.root_dir / split
+
+        self.samples = self._build_sample_index()
+
+        self.rgb_transform = transforms.Compose([
+            transforms.Resize(img_size),
+            transforms.ToTensor(),
+        ])
+        self.mask_transform = transforms.Compose([
+            transforms.Resize(img_size),
+            transforms.ToTensor(),
+        ])
+
+    def _build_sample_index(self):
+        samples = []
+        for scene_dir in self.root_dir.iterdir():
+            if not scene_dir.is_dir():
+                continue
+            for camera_dir in scene_dir.iterdir():
+                if not camera_dir.name.startswith('camera_'):
+                    continue
+
+                rgba_paths = sorted(camera_dir.glob('rgba_*.png'))
+                seg_paths = sorted(camera_dir.glob('segmentation_*.png'))
+
+                for obj_dir in camera_dir.iterdir():
+                    if not obj_dir.name.startswith('obj_'):
+                        continue
+
+                    amodal_paths = sorted(obj_dir.glob('segmentation_*.png'))
+                    amodal_rgb_paths = sorted(obj_dir.glob('rgba_*.png'))
+
+                    if not (len(rgba_paths) == len(seg_paths) == len(amodal_paths) == len(amodal_rgb_paths)):
+                        continue
+
+                    for rgba_path, seg_path, amodal_path, amodal_rgb_path in zip(
+                        rgba_paths, seg_paths, amodal_paths, amodal_rgb_paths
+                    ):
+                        samples.append({
+                            'rgb_path': rgba_path,
+                            'seg_path': seg_path,
+                            'amodal_path': amodal_path,
+                            'amodal_rgb_path': amodal_rgb_path,
+                            'object_id': int(obj_dir.name.split('_')[1]),
+                            'scene': scene_dir.name,
+                            'camera': camera_dir.name
+                        })
+
+        # Limit dataset size if specified
+        if self.max_samples is not None and len(samples) > self.max_samples:
+            # Randomly sample to get diverse examples
+            import random
+            random.seed(42)  # For reproducibility
+            samples = random.sample(samples, self.max_samples)
+            print(f"Dataset limited to {len(samples)} samples")
+
+        # Create train/val split if using validation from train
+        if self.use_val_from_train:
+            import random
+            random.seed(42)  # Ensure reproducible splits
+            random.shuffle(samples)
+
+            val_size = int(len(samples) * self.val_split)
+            if self.split == 'train':
+                samples = samples[val_size:]  # Use remaining samples for training
+                print(f"Train split: {len(samples)} samples")
+            elif self.split == 'val':
+                samples = samples[:val_size]  # Use first samples for validation
+                print(f"Validation split: {len(samples)} samples")
+
+        return samples
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        sample = self.samples[idx]
+
+        # Load images
+        rgb = Image.open(sample['rgb_path']).convert('RGB')
+        seg_map = np.array(Image.open(sample['seg_path']))
+        amodal_mask_img = Image.open(sample['amodal_path']).convert('L')
+        amodal_rgb = Image.open(sample['amodal_rgb_path']).convert('RGB')
+
+        # Compute modal mask (visible part)
+        modal_mask_np = (seg_map == sample['object_id']).astype(np.uint8) * 255
+        modal_mask_img = Image.fromarray(modal_mask_np, mode='L')
+
+        # Transform images and masks
+        rgb = self.rgb_transform(rgb)
+        modal_mask = self.mask_transform(modal_mask_img)
+        amodal_mask = self.mask_transform(amodal_mask_img)
+        amodal_rgb = self.rgb_transform(amodal_rgb)
+
+        # Create occluded mask (parts that are hidden)
+        occluded_mask = amodal_mask - modal_mask
+        occluded_mask = torch.clamp(occluded_mask, 0, 1)
+
+        return {
+            'rgb': rgb,
+            'modal_mask': modal_mask,
+            'amodal_mask': amodal_mask,
+            'occluded_mask': occluded_mask,
+            'amodal_rgb': amodal_rgb,
+        }
+
+
+class ImprovedUNet(nn.Module):
+
+    def __init__(self, in_channels=5, out_channels=3):  # RGB + modal_mask + amodal_mask
+        super().__init__()
+
+        def conv_block(in_ch, out_ch, dropout=0.1):
+            return nn.Sequential(
+                nn.Conv2d(in_ch, out_ch, 3, padding=1),
+                nn.BatchNorm2d(out_ch),
+                nn.ReLU(inplace=True),
+                nn.Dropout2d(dropout),
+                nn.Conv2d(out_ch, out_ch, 3, padding=1),
+                nn.BatchNorm2d(out_ch),
+                nn.ReLU(inplace=True)
+            )
+
+        # Encoder
+        self.down1 = conv_block(in_channels, 64)
+        self.pool1 = nn.MaxPool2d(2)
+        self.down2 = conv_block(64, 128)
+        self.pool2 = nn.MaxPool2d(2)
+        self.down3 = conv_block(128, 256)
+        self.pool3 = nn.MaxPool2d(2)
+        self.down4 = conv_block(256, 512)
+        self.pool4 = nn.MaxPool2d(2)
+
+        # Bottleneck
+        self.middle = conv_block(512, 1024, dropout=0.2)
+
+        # Decoder
+        self.up1 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
+        self.up_block1 = conv_block(1024, 512)
+        self.up2 = nn.ConvTranspose2d(512, 256, 2, stride=2)
+        self.up_block2 = conv_block(512, 256)
+        self.up3 = nn.ConvTranspose2d(256, 128, 2, stride=2)
+        self.up_block3 = conv_block(256, 128)
+        self.up4 = nn.ConvTranspose2d(128, 64, 2, stride=2)
+        self.up_block4 = conv_block(128, 64)
+
+        self.final = nn.Conv2d(64, out_channels, 1)
+
+    def forward(self, x):
+        # Encoder
+        d1 = self.down1(x)
+        d2 = self.down2(self.pool1(d1))
+        d3 = self.down3(self.pool2(d2))
+        d4 = self.down4(self.pool3(d3))
+
+        # Bottleneck
+        m = self.middle(self.pool4(d4))
+
+        # Decoder with skip connections
+        u1 = self.up_block1(torch.cat([self.up1(m), d4], dim=1))
+        u2 = self.up_block2(torch.cat([self.up2(u1), d3], dim=1))
+        u3 = self.up_block3(torch.cat([self.up3(u2), d2], dim=1))
+        u4 = self.up_block4(torch.cat([self.up4(u3), d1], dim=1))
+
+        return torch.sigmoid(self.final(u4))  # Ensure output is in [0,1]
+
+class AmodalCompletionLoss(nn.Module):
+    """Loss that only considers object regions (ignores background)"""
+
+    def __init__(self, occluded_weight=5.0, visible_weight=1.0):
+        super().__init__()
+        self.occluded_weight = occluded_weight
+        self.visible_weight = visible_weight
+        self.lpips_model = lpips.LPIPS(net='alex')
+
+    def forward(self, pred, target, modal_mask, occluded_mask, amodal_mask):
+        # Only compute loss within the amodal mask (object region)
+        device = pred.device
+        self.lpips_model = self.lpips_model.to(device)
+
+        pred_masked = pred * amodal_mask
+        target_masked = target * amodal_mask
+
+
+
+        # Loss on visible parts (within object)
+        visible_region = modal_mask * amodal_mask
+        if visible_region.sum() > 0:
+            visible_loss = F.mse_loss(pred_masked * visible_region, target_masked * visible_region)
+        else:
+            visible_loss = torch.tensor(0.0).to(pred.device)
+
+        # Loss on occluded parts (within object)
+        occluded_region = occluded_mask * amodal_mask
+        if occluded_region.sum() > 0:
+            occluded_loss = F.mse_loss(pred_masked * occluded_region, target_masked * occluded_region)
+        else:
+            occluded_loss = torch.tensor(0.0).to(pred.device)
+
+
+        perceptual_loss = self.lpips_model(pred_masked, target_masked).mean()
+
+        # Boundary consistency within object
+        boundary_mask = F.conv2d(amodal_mask, torch.ones(1,1,3,3).to(amodal_mask.device), padding=1)
+        boundary_mask = ((boundary_mask > 0) & (boundary_mask < 9)).float()
+        boundary_loss = F.mse_loss(pred_masked * boundary_mask, target_masked * boundary_mask)
+
+        total_loss = (self.visible_weight * visible_loss +
+                     self.occluded_weight * occluded_loss +
+                     2.0 * boundary_loss)
+
+        return total_loss, visible_loss, occluded_loss, boundary_loss
+
+
+def train_improved(model, dataloader, optimizer, device, num_epochs):
+    model.train()
+    criterion = AmodalCompletionLoss()
+
+    for epoch in range(num_epochs):
+        total_loss = 0
+        for i, batch in enumerate(dataloader):
+            rgb = batch['rgb'].to(device)
+            modal_mask = batch['modal_mask'].to(device)
+            amodal_mask = batch['amodal_mask'].to(device)
+            occluded_mask = batch['occluded_mask'].to(device)
+            gt_amodal_rgb = batch['amodal_rgb'].to(device)
+
+            input_tensor = torch.cat([rgb, modal_mask, amodal_mask], dim=1)
+
+            optimizer.zero_grad()
+            pred = model(input_tensor)
+
+            loss, vis_loss, occ_loss, boundary_loss = criterion(
+                pred, gt_amodal_rgb, modal_mask, occluded_mask, amodal_mask
+            )
+
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+
+            if i % 16 == 0:
+                print(f"Epoch [{epoch}/{num_epochs}] [{i}/{len(dataloader)}] "
+                      f"Total: {loss.item():.4f}, Visible: {vis_loss.item():.4f}, "
+                      f"Occluded: {occ_loss.item():.4f}, Boundary: {boundary_loss.item():.4f}")
+
+        print(f"Epoch {epoch} Average Loss: {total_loss/len(dataloader):.4f}")
+
+# Usage
+if __name__ == "__main__":
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # Dataset and DataLoader - REDUCED SIZE FOR FASTER TRAINING
+    data_root = "data"
+
+    # Create train dataset (80% of train folder)
+    train_dataset = ModalAmodalDataset(
+        root_dir=data_root,
+        split='train',
+        img_size=(128, 128),
+        max_samples=1000,  # Only use 1000 samples total before split
+        val_split=0.2,     # 20% for validation
+        use_val_from_train=True  # Create val split from train folder
+    )
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=16,
+        shuffle=True,
+        num_workers=2,
+        pin_memory=True,
+        drop_last=True
+    )
+
+    # Create validation dataset (20% of train folder)
+    val_dataset = ModalAmodalDataset(
+        root_dir=data_root,
+        split='val',
+        img_size=(128, 128),
+        max_samples=1000,  # Same max_samples to ensure proper split
+        val_split=0.2,
+        use_val_from_train=True  # Create val split from train folder
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=4,
+        shuffle=True,
+        num_workers=2,
+        pin_memory=True
+    )
+
+    print(f"Training on {len(train_dataset)} samples, {len(train_loader)} batches per epoch")
+    print(f"Validation on {len(val_dataset)} samples, {len(val_loader)} batches")
+
+
+
+
+    model = ImprovedUNet().to(device)
+    model.load_state_dict(torch.load('amodal_completion_model.pth', map_location=device))
+
+
+
+
+
+
+    # Model and optimizer
+    model = model.to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
+
+    # Training
+    #train_improved(model, train_loader, optimizer, device, num_epochs=10)
+
+    # Evaluation and Visualization
+    print("\n" + "="*50)
+    print("EVALUATION RESULTS")
+    print("="*50)
+
+    # Compute metrics
+    metrics = evaluate_metrics(model, val_loader, device)
+    print(f"Overall MSE: {metrics['total_mse']:.6f}")
+    print(f"Occluded Region MSE: {metrics['occluded_mse']:.6f}")
+    print(f"Visible Region MSE: {metrics['visible_mse']:.6f}")
+    print(f"Occluded/Visible MSE Ratio: {metrics['occluded_mse']/metrics['visible_mse']:.2f}")
+
+    # Visualize results
+    print("\nGenerating visualizations...")
+    visualize_results(model, val_loader, device, num_samples=8)
+
+    # Compute metrics
+    image_metrics = calculate_metrics(model, val_loader, device)
+    print(f"PSNR: {image_metrics['psnr']:.4f}")
+    print(f"SSIM: {image_metrics['ssim']:.4f}")
+    print(f"LPIPS: {image_metrics['lpips']:.4f}")
+    print(f"mIoU (pred vs GT): {image_metrics['miou']:.4f}")
+
+# Dataset and DataLoader - REDUCED SIZE FOR FASTER TRAINING
+data_root = "data"
+
+# Create train dataset (80% of train folder)
+train_dataset = ModalAmodalDataset(
+    root_dir=data_root,
+    split='train',
+    img_size=(128, 128),
+    max_samples=1000,  # Only use 1000 samples total before split
+    val_split=0.2,     # 20% for validation
+    use_val_from_train=True  # Create val split from train folder
+)
+train_loader = DataLoader(
+    train_dataset,
+    batch_size=16,
+    shuffle=True,
+    num_workers=2,
+    pin_memory=True,
+    drop_last=True
+)
+
+# Create validation dataset (20% of train folder)
+val_dataset = ModalAmodalDataset(
+    root_dir=data_root,
+    split='val',
+    img_size=(128, 128),
+    max_samples=1000,  # Same max_samples to ensure proper split
+    val_split=0.2,
+    use_val_from_train=True  # Create val split from train folder
+)
+val_loader = DataLoader(
+    val_dataset,
+    batch_size=4,
+    shuffle=True,
+    num_workers=2,
+    pin_memory=True
+)
+
+# Optional: Save model
+torch.save(model.state_dict(), 'amodal_completion_model.pth')
+
+# Evaluation and Visualization
+
+test_dataset = ModalAmodalDataset(
+        root_dir=data_root,
+        split='test',
+        img_size=(128, 128),
+        max_samples=2000  # Only use 1000 samples total before split
+    )
+test_loader = DataLoader(
+        test_dataset,
+        batch_size=8,
+        shuffle=True,
+        num_workers=2,
+        pin_memory=True,
+        drop_last=True
+    )
+
+print("EVALUATION RESULTS")
+print("="*50)
+
+# Compute metrics
+metrics = evaluate_metrics(model, test_loader, device)
+print(f"Overall MSE: {metrics['total_mse']:.6f}")
+print(f"Occluded Region MSE: {metrics['occluded_mse']:.6f}")
+print(f"Visible Region MSE: {metrics['visible_mse']:.6f}")
+print(f"Occluded/Visible MSE Ratio: {metrics['occluded_mse']/metrics['visible_mse']:.2f}")
+
+# Visualize results
+print("\nGenerating visualizations...")
+visualize_results(model, test_loader, device, num_samples=16)
+
+from google.colab import runtime
+runtime.unassign()
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = ImprovedUNet()  # replace with actual class name
+torch.load('amodal_completion_model.pth', map_location=torch.device('cpu'))
+model.to(device)
+model.eval()
+
+# Evaluation and Visualization
+print("\n" + "="*50)
+print("EVALUATION RESULTS")
+print("="*50)
+
+# Compute metrics
+metrics = evaluate_metrics(model, val_loader, device)
+print(f"Overall MSE: {metrics['total_mse']:.6f}")
+print(f"Occluded Region MSE: {metrics['occluded_mse']:.6f}")
+print(f"Visible Region MSE: {metrics['visible_mse']:.6f}")
+print(f"Occluded/Visible MSE Ratio: {metrics['occluded_mse']/metrics['visible_mse']:.2f}")
+
+# Visualize results
+print("\nGenerating visualizations...")
+visualize_results(model, val_loader, device, num_samples=8)
+
+# Compute metrics
+image_metrics = calculate_metrics(model, val_loader, device)
+print(f"PSNR: {image_metrics['psnr']:.4f}")
+print(f"SSIM: {image_metrics['ssim']:.4f}")
+print(f"LPIPS: {image_metrics['lpips']:.4f}")
+print(f"mIoU (pred vs GT): {image_metrics['miou']:.4f}")
+
+model = ImprovedUNet()
+model.eval()