make_truncated_test_set_JWSTex.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efaf81ad-edf6-4c0e-85e5-c63e37ef5a7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.io import fits\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import colors\n",
    "from astropy.visualization import make_lupton_rgb\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from astropy.nddata import Cutout2D\n",
    "from astropy.wcs import WCS\n",
    "\n",
    "import cv2\n",
    "from detectron2.structures import BoxMode\n",
    "from astropy.table import Table\n",
    "import glob\n",
    "from astropy.coordinates import SkyCoord  # High-level coordinates\n",
    "from detectron2.config import LazyConfig, get_cfg, instantiate\n",
    "import os\n",
    "import scipy.stats as stats\n",
    "\n",
    "import astropy.units as u\n",
    "from astropy.coordinates import SkyCoord\n",
    "import h5py\n",
    "import json\n",
    "from detectron2.data import detection_utils as utils\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f55969f6-a1a3-4baf-9a91-490d3869a38e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#Getting all the image data\n",
    "\n",
    "files = sorted(glob.glob('/home/shared/hsc/JWST/images/hlsp_jades_jwst_nircam_goods-s-deep_f*'))\n",
    "JADES_filters = ['090','115','150','200','277','335','356','410','444']\n",
    "\n",
    "JADES_files = []\n",
    "for filt in JADES_filters:\n",
    "    for file in files:\n",
    "        if filt in file:\n",
    "            JADES_files.append(file)\n",
    "\n",
    "            \n",
    "allims = []\n",
    "for file in files:\n",
    "    with fits.open(file) as hdul:\n",
    "        im = hdul[1].data\n",
    "        wcs = WCS(hdul[1].header)\n",
    "        allims.append(im)\n",
    "        \n",
    "        \n",
    "\n",
    "JADESims = []\n",
    "for file in JADES_files:\n",
    "    with fits.open(file) as hdul:\n",
    "        im = hdul[1].data\n",
    "        wcs = WCS(hdul[1].header)\n",
    "        JADESims.append(im)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7879e00-49e7-42f3-b6b1-ce5df89599c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Gets the test metadata (here it's stored in hdf5 format)\n",
    "\n",
    "testfile = '/home/shared/hsc/JWST/processed_data/by_area/flattened_images_test_compresample.hdf5'\n",
    "test_metadatafile = '/home/shared/hsc/JWST/processed_data/by_area/test_compresample_metadata_wcs.hdf5'\n",
    "\n",
    "with h5py.File(test_metadatafile, 'r') as f:\n",
    "    test_metadata = f['metadata_dicts'][:]\n",
    "\n",
    "with h5py.File(testfile, 'r') as fim:\n",
    "    test_images = fim['images'][:]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99fc3bb9-64dc-4295-9968-f67335e2ea4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#determine truncated object IDs\n",
    "\n",
    "def outside_box(box,shape):\n",
    "    if box[0]<0:\n",
    "        return True\n",
    "    elif box[1]<0:\n",
    "        return True\n",
    "    elif box[2]>shape[1]:\n",
    "        return True\n",
    "    elif box[3]>shape[0]:\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "truncated = []\n",
    "for i in range(len(test_metadata)):\n",
    "    d = json.loads(test_metadata[i])\n",
    "    shape=(d['height'],d['width'])\n",
    "    boxes = utils.annotations_to_instances(d['annotations'],shape).gt_boxes.tensor.cpu().numpy()\n",
    "    for j,box in enumerate(boxes):\n",
    "        outside = outside_box(box,shape)\n",
    "        if d['annotations'][j]['redshift']!=-1 and outside:\n",
    "            truncated.append((i,j,d['annotations'][j]['obj_id']))\n",
    "            \n",
    "truncated=np.array(truncated)\n",
    "truncated_ids = np.unique(truncated[:,2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15547ce1-520e-4a25-afc7-0498ce6cec89",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#Get the catalog and divvy into the associated training/test set \n",
    "\n",
    "fncat = '/home/shared/hsc/JWST/catalogs/hlsp_jades_jwst_nircam_goods-s-deep_photometry_v2.0_catalog.fits'\n",
    "fnspecz = '/home/shared/hsc/JWST/catalogs/JADES_GOODS_zspec_cleaned.fits'\n",
    "fndect = '/home/shared/hsc/JWST/images/hlsp_jades_jwst_nircam_goods-s-deep_detection_v2.0_drz.fits'\n",
    "fnseg = '/home/shared/hsc/JWST/images/hlsp_jades_jwst_nircam_goods-s-deep_segmentation_v2.0_drz.fits'\n",
    "\n",
    "dphot = Table.read(fncat, hdu=2).to_pandas()\n",
    "dspecz = Table.read(fnspecz,hdu=1).to_pandas()\n",
    "dseg = Table.read(fncat, hdu=3).to_pandas()\n",
    "\n",
    "\n",
    "\n",
    "with fits.open(fnseg) as hdul:\n",
    "    hdul.info()\n",
    "    segim = hdul[0].data\n",
    "    segh = hdul[0].header\n",
    "    \n",
    "\n",
    "with fits.open(fndect) as hdul:\n",
    "    hdul.info()\n",
    "    dectim = hdul[0].data\n",
    "    dect = hdul[0].header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7a6f950-fa4d-464b-928b-ed54172ae098",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#Annotation function\n",
    "\n",
    "def get_metadata(segcutout, dspecz, imid):\n",
    "    segimcut = segcutout.data\n",
    "    d = {}\n",
    "    annos = []\n",
    "    for s in np.unique(segimcut):\n",
    "        mask = np.zeros(segimcut.shape)\n",
    "        if s == 0:\n",
    "            continue\n",
    "        s0i = np.where(segimcut == s)\n",
    "        mask[s0i] =1\n",
    "\n",
    "        #x0 = s0i[1].min()\n",
    "        #x1 = s0i[1].max()\n",
    "        #y0 = s0i[0].min()\n",
    "        #y1 = s0i[0].max()\n",
    "\n",
    "        #h = int(y1-y0)\n",
    "        #w = int(x1-x0)\n",
    "        \n",
    "     \n",
    "        x0,y0 = segcutout.to_cutout_position((dseg[dseg['ID']==s]['BBOX_XMIN'].values[0],dseg[dseg['ID']==s]['BBOX_YMIN'].values[0]))\n",
    "        x1,y1 = segcutout.to_cutout_position((dseg[dseg['ID']==s]['BBOX_XMAX'].values[0],dseg[dseg['ID']==s]['BBOX_YMAX'].values[0]))\n",
    "\n",
    "        h = int(y1-y0)\n",
    "        w = int(x1-x0)\n",
    "\n",
    "        contours, hierarchy = cv2.findContours(\n",
    "                (mask).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE\n",
    "            )\n",
    "\n",
    "        segmentation = []\n",
    "        for contour in contours:\n",
    "            # contour = [x1, y1, ..., xn, yn]\n",
    "            contour = contour.flatten()\n",
    "            if len(contour) > 4:\n",
    "                segmentation.append(contour.tolist())\n",
    "        \n",
    "        specmatch = np.where(dspecz.ID==s)[0]\n",
    "        if len(specmatch) ==0:\n",
    "            redshift = -1\n",
    "        else:\n",
    "            #redshift = dspecz.iloc[specmatch]['z'].value\n",
    "            redshift = dspecz['z'].values[specmatch[0]]\n",
    "            \n",
    "        obj = {\n",
    "            \"bbox\": [int(x0), int(y0), w, h],\n",
    "            \"area\": w * h,\n",
    "            \"bbox_mode\": BoxMode.XYWH_ABS,\n",
    "            \"segmentation\": segmentation,\n",
    "            \"category_id\": 0,\n",
    "            \"obj_id\": int(s),\n",
    "            \"redshift\":redshift\n",
    "        }\n",
    "        \n",
    "        annos.append(obj)\n",
    "\n",
    "    height, width = mask.shape\n",
    "\n",
    "\n",
    "\n",
    "    d[\"annotations\"] = annos\n",
    "    d['height'] = height\n",
    "    d['width'] = width\n",
    "    d[\"image_id\"] = imid\n",
    "    \n",
    "    return d\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68914d20-cf2d-4dc6-bb36-664e844dd673",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#Example of making a cutout centered on the jth truncated object\n",
    "#Get the RA and DEC and then cutout from the main image using its WCS\n",
    "\n",
    "j=38\n",
    "coords = SkyCoord(ra=dspecz.RA.values[np.where(dspecz.ID.values==trunc_ids[j])]*u.degree, dec=dspecz.DEC.values[np.where(dspecz.ID.values==trunc_ids[j])]*u.degree)\n",
    "\n",
    "image = []\n",
    "for im in JADESims:\n",
    "    cutout = Cutout2D(im, position=coords, size=(512,512), wcs=wcs)\n",
    "    image.append(cutout.data)\n",
    "image = np.array(image)\n",
    "\n",
    "\n",
    "segcutout = Cutout2D(segim, position=coords, size=(512,512), wcs=wcs)\n",
    "#segimcut = segcutout.data\n",
    "\n",
    "d = get_metadata(segcutout, dspecz, j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dff3cd7-c449-4a59-8a32-61a963ab77e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from detectron2.data import MetadataCatalog, DatasetCatalog\n",
    "\n",
    "astrotest_metadata = MetadataCatalog.get(\"astro_test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "176072bf-d15b-4270-8213-c5f88d816657",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from deepdisc.astrodet.visualizer import Visualizer\n",
    "from deepdisc.astrodet.visualizer import ColorMode\n",
    "\n",
    "from astropy.visualization import make_lupton_rgb\n",
    "\n",
    "plt.figure(figsize=(14,7))\n",
    "b1 = image[2]\n",
    "b2 = image[1]\n",
    "b3 = image[0]\n",
    "\n",
    "img = make_lupton_rgb(b1, b2, b3, minimum=0, stretch=0.1, Q=5)\n",
    "\n",
    "print(\"total instances:\", len(d[\"annotations\"]))\n",
    "v0 = Visualizer(\n",
    "    img,\n",
    "    metadata=astrotest_metadata,\n",
    "    scale=5,\n",
    "    instance_mode=ColorMode.SEGMENTATION,  # remove the colors of unsegmented pixels. This option is only available for segmentation models\n",
    ")\n",
    "groundTruth = v0.draw_dataset_dict(d, lf=False, alpha=0.01)\n",
    "\n",
    "ax1 = plt.subplot(1, 1, 1)\n",
    "ax1.imshow(groundTruth.get_image())\n",
    "ax1.axis(\"off\")\n",
    "ax1.set_title('Ground Truth')\n",
    "#ax2 = plt.subplot(1, 2, 2)\n",
    "#ax2.imshow(image[0],norm=colors.LogNorm())\n",
    "#ax2.axis(\"off\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c125a3b-5e6a-49ce-a4ed-c0dedf128441",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "#Loops over the truncated objects\n",
    "\n",
    "dicts = []\n",
    "images = []\n",
    "for j in range(len(truncated_ids)):\n",
    "    print(j)\n",
    "    \n",
    "    coords = SkyCoord(ra=dspecz.RA.values[np.where(dspecz.ID.values==truncated_ids[j])]*u.degree, dec=dspecz.DEC.values[np.where(dspecz.ID.values==truncated_ids[j])]*u.degree)\n",
    "\n",
    "    \n",
    "    image = []\n",
    "    for im in JADESims:\n",
    "        cutout = Cutout2D(im, position=coords, size=(512,512), wcs=wcs)\n",
    "        image.append(cutout.data)\n",
    "    image = np.array(image)\n",
    "    \n",
    "    \n",
    "    segcutout = Cutout2D(segim, position=coords, size=(512,512), wcs=wcs)\n",
    "    #segimcut = segcutout.data\n",
    "    \n",
    "    d = get_metadata(segcutout, dspecz, j)\n",
    "    if np.all(image==0):\n",
    "        continue\n",
    "        \n",
    "    ddict={}\n",
    "    ddict['wcs'] = cutout.wcs.to_header_string()\n",
    "    ddict['height'] = image.shape[1]\n",
    "    ddict['width'] = image.shape[2]\n",
    "    ddict['annotations'] = d['annotations']\n",
    "    ddict['image_id'] = int(truncated_ids[j])\n",
    "\n",
    "    dicts.append(ddict)\n",
    "    images.append(image.flatten())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a28a374-fb38-4471-ac44-dc131b4fc047",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#Saves in HDF5 format (for RAIL photozs)\n",
    "\n",
    "import deepdisc.data_format.conversions as conversions\n",
    "\n",
    "trunc_images = np.array(images)\n",
    "\n",
    "with h5py.File('/home/shared/hsc/JWST/processed_data/by_area/flattened_images_test_compresample_trunc.hdf5', \"w\") as f:\n",
    "    data = f.create_dataset(\"images\", data=trunc_images)\n",
    "    \n",
    "conversions.ddict_to_hdf5(dicts,'/home/shared/hsc/JWST/processed_data/by_area/test_comp_resample_trunc_metadata_wcs.hdf5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c902e92a-d123-4e3d-8423-d54d8d68154a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ef74fd6-54f5-4378-bd8b-9e45a4eddf70",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-ddrailnv]",
   "language": "python",
   "name": "conda-env-.conda-ddrailnv-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}