dataloader.py

import os
import cv2
import json
import torch
import torchvision
import numpy as np
import pandas as pd
import torch.nn as nn
import albumentations as A
import matplotlib.pyplot as plt
import torch.nn.functional as F

from PIL import Image
from tqdm import tqdm
from collections import defaultdict
from albumentations.pytorch import ToTensorV2
from torch.utils.data import Dataset, DataLoader
from torchvision.models.detection import MaskRCNN
from sklearn.model_selection import train_test_split
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.transforms import functional as F_transforms

from torchvision.ops import box_iou
from torchvision.models.detection.image_list import ImageList

from torchvision.models.detection import maskrcnn_resnet50_fpn


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

@torch.no_grad()
def run_rpn_only(images):
    """
    images: list of tensors [C,H,W], already normalized
    returns: list of proposal tensors [N,4] per image
    """
    rpn_model = maskrcnn_resnet50_fpn(weights="DEFAULT")
    rpn_model.eval()
    rpn_model.to(device)

    # Freeze everything
    for p in rpn_model.parameters():
        p.requires_grad = False

    # Backbone
    features = rpn_model.backbone(images)

    if isinstance(features, torch.Tensor):
        features = {"0": features}

    # Image sizes
    image_sizes = [img.shape[-2:] for img in images]
    image_list = ImageList(images, image_sizes)

    # RPN proposals
    proposals, _ = rpn_model.rpn(image_list, features)
    return proposals

from torchvision.ops import box_iou

def select_bg_boxes_from_rpn(rpn_boxes, gt_boxes, max_bg=3, iou_thr=0.1):
    """
    rpn_boxes: [P,4]
    gt_boxes:  [G,4]
    """
    if len(gt_boxes) == 0:
        return rpn_boxes[:max_bg]

    ious = box_iou(rpn_boxes, gt_boxes)   # [P,G]
    max_iou, _ = ious.max(dim=1)

    bg_boxes = rpn_boxes[max_iou < iou_thr]
    return bg_boxes[:max_bg]

def resize_mask(combined_mask, target_image):
    """
    Resizes a mask to match target_image size.
    Accepts combined_mask of shape [1,H,W] or [1,1,H,W].
    """
    # Ensure mask has shape [N, C, H, W] for F.interpolate
    if combined_mask.ndim == 3:
        combined_mask = combined_mask.unsqueeze(0)  # -> [1, 1, H, W]
    elif combined_mask.ndim != 4:
        raise ValueError(f"Unexpected mask shape: {combined_mask.shape}")

    # Get target height and width
    if isinstance(target_image, torch.Tensor):
        if target_image.ndim == 3:  # (C,H,W) or (H,W,C)
            if target_image.shape[0] in (1, 3):      # (C,H,W)
                H_img, W_img = target_image.shape[1], target_image.shape[2]
            else:                                    # (H,W,C)
                H_img, W_img = target_image.shape[0], target_image.shape[1]

        elif target_image.ndim == 2:  # (H,W)
            H_img, W_img = target_image.shape

        else:
            raise ValueError(f"Unexpected target_image shape: {target_image.shape}")
    else:  # assume numpy
        H_img, W_img = target_image.shape[:2]

    # Interpolate mask to target size
    mask_resized = F.interpolate(
        combined_mask.float(),
        size=(H_img, W_img),
        mode='nearest'
    )
    return mask_resized

def combine_resize_submasks(output, target_image, pixel_thresh=None, mask_thresh=None, already_binary=False):

    if len(output["masks"]) == 0:
        return torch.zeros(
            target_image.shape[0],
            target_image.shape[1],
            device=target_image.device
        )

    masks = output["masks"].squeeze(1)  # [N,H,W]
    combined = torch.zeros_like(masks[0], dtype=torch.float32)
    mask_confs = []
    submasks_above_thresh = 0

    if already_binary: # Target was passed, submasks are dtype=uint8
        for j in range(len(masks)):
            combined += masks[j].float()
            mask_confs.append(1.0)
    else:
        for j in range(len(masks)):
            mask_probs = masks[j]      # Pixel forgery prob in E (0, 1)
            binary_mask = mask_probs > pixel_thresh      # Thresholded gives binary
            """ So mask_probs is soft, binary_mask is smaller """
            #print(masks[j].max(), masks[j].min(), masks[j].float().mean())

            if binary_mask.sum() == 0:
                continue

            """ Mean of the probs inside the binary mask E [0,1] """
            mask_conf = mask_probs[binary_mask].float().mean()
            mask_confs.append(mask_conf)

            """ Combine only binaries with high forgery probability """
            if mask_conf > mask_thresh:
                combined += binary_mask.float()
                submasks_above_thresh += 1

    #print(f"Submasks above thresh: {submasks_above_thresh}")
    combined = combined.unsqueeze(0).unsqueeze(0)   # [1,1,Hm,Wm]
    combined = resize_mask(combined, target_image)

    return combined.squeeze(0).squeeze(0), torch.tensor(mask_confs)


class ForgeryDataset(Dataset):
    def __init__(self, authentic_path, forged_path, masks_path, transform=None, is_train=True):
        self.transform = transform
        self.is_train = is_train

        # Collect all data samples
        self.samples = []

        # Forged images
        for file in sorted(os.listdir(forged_path)):
            if file[0] == ".":
                continue
            img_path = os.path.join(forged_path, file)
            base_name = file.split('.')[0]
            mask_path = os.path.join(masks_path, f"{base_name}.npy")

            self.samples.append({
                'image_path': img_path,
                'mask_path': mask_path,
                'is_forged': True,
                'image_id': base_name
            })

        # Authentic images
        if (authentic_path is not None):
            for file in sorted(os.listdir(authentic_path)):
                if file[0] == ".":
                    continue
                img_path = os.path.join(authentic_path, file)
                base_name = file.split('.')[0]
                mask_path = os.path.join(masks_path, f"{base_name}.npy")

                self.samples.append({
                    'image_path': img_path,
                    'mask_path': mask_path,
                    'is_forged': False,
                    'image_id': base_name
                })

    def __len__(self):
        return len(self.samples)

    def get_raw_img_mask(self, idx):
        sample = self.samples[idx]
        image_raw = Image.open(sample['image_path']).convert('RGB')
        image_raw = np.array(image_raw)  # (H, W, 3)

        if os.path.exists(sample['mask_path']):
            mask = np.load(sample['mask_path'])

            return image_raw, mask
        else:
            return image_raw, np.asarray([0])

    def get_image_props(self, image, mask):
        boxes, labels, masks = self.mask_to_boxes(image, mask)
        mask_np = masks.cpu().numpy() if isinstance(masks, torch.Tensor) else masks
        mask_whiteness = mask_np.sum() / (image.shape[1] * image.shape[2])

        return {
            "Npixels" : len(image[0])*len(image[1]),
            "WhiteNess" : mask_whiteness
            }

    def get_filename(self, idx):
        return self.samples[idx]

    def __getitem__(self, idx):
        sample = self.samples[idx]

        # Load image
        image = Image.open(sample['image_path']).convert('RGB')
        image = np.array(image)  # (H, W, 3)

        # Load and process mask
        if os.path.exists(sample['mask_path']):
            mask = np.load(sample['mask_path'])

            # Handle multi-channel masks
            if mask.ndim == 3:
                if mask.shape[0] <= 10:  # channels first (C, H, W)
                    mask = np.any(mask, axis=0)
                elif mask.shape[-1] <= 10:  # channels last (H, W, C)
                    mask = np.any(mask, axis=-1)
                else:
                    raise ValueError(f"Ambiguous 3D mask shape: {mask.shape}")

            mask = (mask > 0).astype(np.uint8)
        else:
            mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)

        # Resize mask to match image if needed
        H_img, W_img = image.shape[:2]
        if mask.shape != (H_img, W_img):
            print(f"[WARN] pre-resizing mask {mask.shape} -> {(H_img, W_img)}")
            mask = cv2.resize(mask.astype(np.uint8), (W_img, H_img), interpolation=cv2.INTER_NEAREST)

        # Apply transformations
        if self.transform:
            transformed = self.transform(image=image, mask=mask)
            image = transformed['image']
            mask = transformed['mask']
        else:
            image = F_transforms.to_tensor(image)
            mask = torch.tensor(mask, dtype=torch.uint8)

        # Prepare targets for Mask R-CNN
        if sample['is_forged'] and mask.sum() > 0:
            boxes, labels, masks = self.mask_to_boxes(image, mask)

            target = {
                'boxes': boxes,
                'labels': labels,
                'masks': masks,
                'image_id': torch.tensor([idx]),
                'area': (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]),
                'iscrowd': torch.zeros((len(boxes),), dtype=torch.int64)
            }
        else:
            # For authentic images or images without masks
            target = {
                'boxes': torch.zeros((0, 4), dtype=torch.float32),
                'labels': torch.zeros(0, dtype=torch.int64),
                'masks': torch.zeros((0, image.shape[1], image.shape[2]), dtype=torch.uint8),
                'image_id': torch.tensor([idx]),
                'area': torch.zeros(0, dtype=torch.float32),
                'iscrowd': torch.zeros((0,), dtype=torch.int64)
            }

        return image, target, self.samples[idx]['image_path'], idx  # return filename too

    def mask_to_boxes(self, image, mask, plot = False):
        """Convert segmentation mask to bounding boxes for Mask R-CNN"""
        if isinstance(mask, torch.Tensor):
            mask_np = mask.detach().cpu().numpy()
        else:
            mask_np = np.array(mask)

        # --- FIX: squeeze extra dimensions ---
        mask_np = np.squeeze(mask_np)
        if mask_np.ndim != 2:
            raise ValueError("mask_to_boxes expects a single 2D mask")

        # --- FIX: ensure binary + correct type ---
        mask_np = (mask_np > 0).astype(np.uint8)

        # Safety: fall back to empty if shape still wrong
        if mask_np.ndim != 2:
            return (torch.zeros((0,4),dtype=torch.float32),
                    torch.zeros((0,),dtype=torch.int64),
                    torch.zeros((0, *mask_np.shape[-2:]),dtype=torch.uint8))

        contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        boxes = []
        masks = []

        for contour in contours:
            if len(contour) > 0:
                x, y, w, h = cv2.boundingRect(contour)
                # Filter out very small regions
                if w > 5 and h > 5:
                    boxes.append([x, y, x + w, y + h])
                    # Create binary mask for this contour
                    contour_mask = np.zeros_like(mask_np)
                    cv2.fillPoly(contour_mask, [contour], 1)
                    masks.append(contour_mask)

        if boxes:
            boxes = torch.tensor(boxes, dtype=torch.float32)
            labels = torch.ones((len(boxes),), dtype=torch.int64)
            masks = torch.tensor(np.array(masks), dtype=torch.uint8)
        else:
            boxes = torch.zeros((0, 4), dtype=torch.float32)
            labels = torch.zeros(0, dtype=torch.int64)
            masks = torch.zeros((0, mask_np.shape[0], mask_np.shape[1]), dtype=torch.uint8)

        return boxes, labels, masks

base_path = "../recodai-luc-scientific-image-forgery-detection/"
paths = {
        'train_authentic': os.path.join(base_path, "train_images/authentic"),
        'train_forged': os.path.join(base_path, "train_images/forged"),
        'train_masks': os.path.join(base_path, "train_masks"),
        'test_images': os.path.join(base_path, "test_images"),
    }

tenimg_paths = {
        'train_authentic': os.path.join(base_path, "train_images/10img/authentic"),
        'train_forged': os.path.join(base_path, "train_images/10img/forged"),
        'train_masks': os.path.join(base_path, "train_images/10img/masks")
    }

# Transformations for learning
train_transform = A.Compose([
    A.Resize(512, 512, interpolation=cv2.INTER_NEAREST),

    A.HorizontalFlip(p=0.5),
    #A.VerticalFlip(p=0.5),
    #A.RandomRotate90(p=0.5),

    # Color / lighting / noise
    A.RandomBrightnessContrast(p=0.3, brightness_limit=0.15, contrast_limit=0.15),
    A.HueSaturationValue(p=0.2),
    A.GaussNoise(p=0.1),


    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
])

val_transform = A.Compose([
    A.Resize(512, 512, interpolation=cv2.INTER_NEAREST),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
])

if __name__ == "__main__":

    print("test")