Task2_create_patch.py

import os
import tifffile
import numpy as np
import json
import random
from PIL import Image
from collections import defaultdict

SEED = 42
random.seed(SEED)
np.random.seed(SEED)
#initializations
PATCH_SIZE = 100
HALF = PATCH_SIZE // 2
BASE_DIR = os.path.expanduser('~/cvcw/Dataset_Splits') #change path to your dataset
OUTPUT_DIR = os.path.expanduser('~/cvcw/task2_patches_fixed')
TRAIN_IMAGE_DIR = os.path.join(BASE_DIR, 'train', 'image')
TRAIN_NUCLEI_DIR = os.path.join(BASE_DIR, 'train', 'nuclei')
VAL_IMAGE_DIR = os.path.join(BASE_DIR, 'validation', 'image')
VAL_NUCLEI_DIR = os.path.join(BASE_DIR, 'validation', 'nuclei')

CLASS_MAP = {
    'nuclei_tumor': 'Tumor',
    'nuclei_lymphocyte': 'Lymphocyte',
    'nuclei_histiocyte': 'Histiocyte',
}
CLASSES = ['Tumor', 'Lymphocyte', 'Histiocyte']

TRAIN_PER_CLASS = 2500
CONTRASTIVE_PER_CLASS = 2500
VAL_PER_CLASS = 700

def ensure_dirs():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    for split_name in ['train', 'contrastive', 'val']:
        for cls in CLASSES:
            os.makedirs(os.path.join(OUTPUT_DIR, split_name, cls), exist_ok=True)

def load_geojson(path):
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)

def get_all_coords(geometry):
    geo_type = geometry.get('type')
    if geo_type == 'Polygon':
        rings = geometry['coordinates']
        coords = []
        for ring in rings:
            coords.extend(ring)
        return np.array(coords, dtype=np.float32)
    if geo_type == 'MultiPolygon':
        coords = []
        for polygon in geometry['coordinates']:
            for ring in polygon:
                coords.extend(ring)
        return np.array(coords, dtype=np.float32)
    return None
def get_center_from_geometry(geometry): #Use bounding-box center instead of centroid
    coords = get_all_coords(geometry)
    if coords is None or len(coords) == 0:
        return None
    xs = coords[:, 0]
    ys = coords[:, 1]
    x_min, x_max = xs.min(), xs.max()
    y_min, y_max = ys.min(), ys.max()
    cx = int(round((x_min + x_max) / 2.0))
    cy = int(round((y_min + y_max) / 2.0))
    return cx, cy

def crop_with_zero_padding(image, cx, cy, patch_size=100): #crop with zeropadding 
    half = patch_size // 2
    h, w = image.shape[:2]
    x1 = cx - half
    x2 = cx + half
    y1 = cy - half
    y2 = cy + half
    patch = np.zeros((patch_size, patch_size, 3), dtype=np.uint8)
    src_x1 = max(0, x1)
    src_x2 = min(w, x2)
    src_y1 = max(0, y1)
    src_y2 = min(h, y2)
    dst_x1 = src_x1 - x1
    dst_x2 = dst_x1 + (src_x2 - src_x1)
    dst_y1 = src_y1 - y1
    dst_y2 = dst_y1 + (src_y2 - src_y1)
    patch[dst_y1:dst_y2, dst_x1:dst_x2] = image[src_y1:src_y2, src_x1:src_x2, :3]
    return patch

def read_slide(path): # ensure consistent 3-channel uint8 format across slides
    image = tifffile.imread(path)
    if image.ndim == 4:
        image = image[0]
    if image.ndim == 2:
        image = np.stack([image] * 3, axis=-1)
    if image.shape[-1] > 3:
        image = image[..., :3]
    if image.dtype != np.uint8:
        image = np.clip(image, 0, 255).astype(np.uint8)
    return image

def collect_cand(image_dir, nuclei_dir, split_label): # extract candidate patches per class from geojson
    candidates = {cls: [] for cls in CLASSES}
    tif_files = sorted([f for f in os.listdir(image_dir) if f.endswith('.tif')])
    print(f"Collecting candidates from {split_label}: {len(tif_files)} slides...")
    for tif_file in tif_files:
        tif_path = os.path.join(image_dir, tif_file)
        geojson_name = tif_file.replace('.tif', '_nuclei.geojson')
        geojson_path = os.path.join(nuclei_dir, geojson_name)
        if not os.path.exists(geojson_path):
            print(f"  Missing GeoJSON for {tif_file} -- skipping")
            continue
        try:
            image = read_slide(tif_path)
            geojson = load_geojson(geojson_path)
        except Exception as e:
            print(f"  Failed to load {tif_file}: {e}")
            continue
        features = geojson.get('features', [])
        kept_this_slide = 0
        for feature_idx, feature in enumerate(features):
            try:
                raw_class = feature['properties']['classification']['name']
            except Exception:
                continue
            if raw_class not in CLASS_MAP:
                continue
            cls_name = CLASS_MAP[raw_class]
            center = get_center_from_geometry(feature.get('geometry', {}))
            if center is None:
                continue
            cx, cy = center
            patch = crop_with_zero_padding(image, cx, cy, patch_size=PATCH_SIZE)
            candidates[cls_name].append({
                'patch': patch,
                'source_split': split_label,
                'slide_name': os.path.splitext(tif_file)[0],
                'feature_idx': int(feature_idx),
                'class_name': cls_name,
                'center_x': int(cx),
                'center_y': int(cy),
                'height': int(image.shape[0]),
                'width': int(image.shape[1]),
            })
            kept_this_slide += 1
        print(f"  {tif_file}: kept {kept_this_slide} nuclei patches")
    for cls in CLASSES:
        print(f"Collected {len(candidates[cls])} total candidates for {cls} from {split_label}")

    return candidates

def patch_rec_save(record, split_name, index_within_split):
    cls_name = record['class_name']
    out_dir = os.path.join(OUTPUT_DIR, split_name, cls_name)
    filename = (
        f"{record['slide_name']}"
        f"_x{record['center_x']}"
        f"_y{record['center_y']}"
        f"_idx{record['feature_idx']}"
        f"_{index_within_split:05d}.npy"
    )
    out_path = os.path.join(out_dir, filename)
    np.save(out_path, record['patch'])
    meta = {
        'file': out_path,
        'split': split_name,
        'class_name': cls_name,
        'source_split': record['source_split'],
        'slide_name': record['slide_name'],
        'feature_idx': record['feature_idx'],
        'center_x': record['center_x'],
        'center_y': record['center_y'],
        'image_height': record['height'],
        'image_width': record['width'],
        'patch_size': PATCH_SIZE,
    }
    return meta
def assign_n_save(train_candidates, val_candidates):
    train_meta = []
    contrastive_meta = []
    val_meta = []
    summary = defaultdict(dict)
    for cls in CLASSES:
        random.shuffle(train_candidates[cls])
        random.shuffle(val_candidates[cls])
        needed_train_total = TRAIN_PER_CLASS + CONTRASTIVE_PER_CLASS # enforce fixed number of samples/class to balance dataset
        available_train = len(train_candidates[cls])
        if available_train < needed_train_total:
            print(f"WARNING: {cls} has only {available_train} train candidates; needed {needed_train_total}")
        train_selected = train_candidates[cls][:TRAIN_PER_CLASS]
        contrastive_selected = train_candidates[cls][TRAIN_PER_CLASS:TRAIN_PER_CLASS + CONTRASTIVE_PER_CLASS]
        val_selected = val_candidates[cls][:VAL_PER_CLASS]
        for i, record in enumerate(train_selected):
            train_meta.append(patch_rec_save(record, 'train', i))
        for i, record in enumerate(contrastive_selected):
            contrastive_meta.append(patch_rec_save(record, 'contrastive', i))
        for i, record in enumerate(val_selected):
            val_meta.append(patch_rec_save(record, 'val', i))
        summary[cls]['train_saved'] = len(train_selected)
        summary[cls]['contrastive_saved'] = len(contrastive_selected)
        summary[cls]['val_saved'] = len(val_selected)
        summary[cls]['train_candidates_available'] = available_train
        summary[cls]['val_candidates_available'] = len(val_candidates[cls])
    return train_meta, contrastive_meta, val_meta, summary

def writejson(path, data):
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2)
if __name__ == '__main__':
    ensure_dirs()
    print('Starting fixed patch extraction...')
    print(f'BASE_DIR: {BASE_DIR}')
    print(f'OUTPUT_DIR: {OUTPUT_DIR}')
    print(f'PATCH_SIZE: {PATCH_SIZE}')
    train_candidates = collect_cand(image_dir=TRAIN_IMAGE_DIR,nuclei_dir=TRAIN_NUCLEI_DIR,split_label='train')
    val_candidates = collect_cand(image_dir=VAL_IMAGE_DIR,nuclei_dir=VAL_NUCLEI_DIR,split_label='validation')
    train_meta, contrastive_meta, val_meta, summary = assign_n_save(train_candidates=train_candidates,val_candidates=val_candidates)
    
    writejson(os.path.join(OUTPUT_DIR, 'train_metadata.json'), train_meta)
    writejson(os.path.join(OUTPUT_DIR, 'contrastive_metadata.json'), contrastive_meta)
    writejson(os.path.join(OUTPUT_DIR, 'val_metadata.json'), val_meta)
    writejson(os.path.join(OUTPUT_DIR, 'summary.json'), summary)
    print('Final Count:')
   
    for cls in CLASSES:
        print(
            f"{cls:12s} , "
            f"train={summary[cls]['train_saved']:4d} , "
            f"contrastive={summary[cls]['contrastive_saved']:4d} , "
            f"val={summary[cls]['val_saved']:4d}"
        )
    print('\nSaved metadata:')
    print(os.path.join(OUTPUT_DIR, 'train_metadata.json'))
    print(os.path.join(OUTPUT_DIR, 'contrastive_metadata.json'))
    print(os.path.join(OUTPUT_DIR, 'val_metadata.json'))
    print(os.path.join(OUTPUT_DIR, 'summary.json'))