text:
""" CS4277/CS5477 Lab 4: Plane Sweep Stereo
See accompanying Jupyter notebook (lab4.ipynb) for instructions.

Name: Shicheng Chen
Email: e0534721@u.nus.edu
Student ID: A0215003A
"""
import json
import os

import cv2
import numpy as np
from scipy.spatial.transform import Rotation
import scipy.ndimage

"""Helper functions: You should not have to touch the following functions.
"""
class Image(object):
    """
    Image class. You might find the following member variables useful:
    - image: RGB image (HxWx3) of dtype np.float64
    - pose_mat: 3x4 Camera extrinsics that transforms points from world to
        camera frame
    """
    def __init__(self, qvec, tvec, name, root_folder=''):
        self.qvec = qvec
        self.tvec = tvec
        self.name = name  # image filename
        self._image = self.load_image(os.path.join(root_folder, name))

        # Extrinsic matrix: Transforms from world to camera frame
        self.pose_mat = self.make_extrinsic(qvec, tvec)

    def __repr__(self):
        return '{}: qvec={}\n tvec={}'.format(
            self.name, self.qvec, self.tvec
        )

    @property
    def image(self):
        return self._image.copy()

    @staticmethod
    def load_image(path):
        """Loads image and converts it to float64"""
        im = cv2.imread(path)
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        return im.astype(np.float64) / 255.0

    @staticmethod
    def make_extrinsic(qvec, tvec):
        """ Make 3x4 camera extrinsic matrix from colmap pose

        Args:
            qvec: Quaternion as per colmap format (q_cv) in the order
                  q_w, q_x, q_y, q_z
            tvec: translation as per colmap format (t_cv)

        Returns:

        """
        rotation = Rotation.from_quat(np.roll(qvec, -1))
        return np.concatenate([rotation.as_dcm(), tvec[:, None]], axis=1)

def write_json(outfile, images, intrinsic_matrix, img_hw):
    """Write metadata to json file.

    Args:
        outfile (str): File to write to
        images (list): List of Images
        intrinsic_matrix (np.ndarray): 3x3 intrinsic matrix
        img_hw (tuple): (image height, image width)
    """

    img_height, img_width = img_hw

    images_meta = []
    for im in images:
        images_meta.append({
            'name': im.name,
            'qvec': im.qvec.tolist(),
            'tvec': im.tvec.tolist(),
        })

    data = {
        'img_height': img_height,
        'img_width': img_width,
        'K': intrinsic_matrix.tolist(),
        'images': images_meta
    }
    with open(outfile, 'w') as fid:
        json.dump(data, fid, indent=2)

def load_data(root_folder):
    """Loads dataset.

    Args:
        root_folder (str): Path to data folder. Should contain metadata.json

    Returns:
        images, K, img_hw
    """
    print('Loading data from {}...'.format(root_folder))
    with open(os.path.join(root_folder, 'metadata.json')) as fid:
        metadata = json.load(fid)

    images = []
    for im in metadata['images']:
        images.append(Image(np.array(im['qvec']), np.array(im['tvec']),
                            im['name'], root_folder=root_folder))
    img_hw = (metadata['img_height'], metadata['img_width'])
    K = np.array(metadata['K'])

    print('Loaded data containing {} images.'.format(len(images)))
    return images, K, img_hw

def rgb2hex(rgb):
    """Converts color representation into hexadecimal representation for K3D

    Args:
        rgb (np.ndarray): (N, 3) array holding colors

    Returns:
        hex (np.ndarray): array (N, ) of size N, each element indicates the
          color, e.g. 0x0000FF = blue
    """
    rgb_uint = (rgb * 255).astype(np.uint8)
    hex = np.sum(rgb_uint * np.array([[256 ** 2, 256, 1]]),
                 axis=1).astype(np.uint32)
    return hex

"""Functions to be implemented
"""
# Part 1
def compute_relative_pose(cam_pose, ref_pose):
    """Compute relative pose between two cameras

     Args:
        cam_pose (np.ndarray): Extrinsic matrix of camera of interest C_i (3,4).
          Transforms points in world frame to camera frame, i.e.
            x_i = C_i @ x_w  (taking into account homogeneous dimensions)
        ref_pose (np.ndarray): Extrinsic matrix of reference camera C_r (3,4)

    Returns:
        relative_pose (np.ndarray): Relative pose of size (3,4). Should transform
          points in C_r to C_i, i.e. x_i = M @ x_r

    Prohibited functions:
        Do NOT use np.linalg.inv() or similar functions
    """
    relative_pose = np.zeros((3, 4), dtype=np.float64)

    """ YOUR CODE STARTS HERE """

    Ri,Rr=cam_pose[:,:-1],ref_pose[:,:-1]
    ti,tr=cam_pose[:,-1:],ref_pose[:,-1:]
    relative_pose[:,:-1]=Ri@Rr.T
    relative_pose[:, -1:]=ti-Ri@Rr.T@tr

    """ YOUR CODE ENDS HERE """
    return relative_pose

def get_plane_sweep_homographies(K, relative_pose, inv_depths):
    """Compute plane sweep homographies, assuming fronto parallel planes w.r.t.
    reference camera

    Args:
        K (np.ndarray): Camera intrinsic matrix (3,3)
        relative_pose (np.ndarray): Relative pose between the two cameras
          of shape (3, 4)
        inv_depths (np.ndarray): Inverse depths to warp of size (D, )

    Returns:
        homographies (D, 3, 3)
    """

    """ YOUR CODE STARTS HERE """
    r=relative_pose[:,:-1]
    t=relative_pose[:,-1:]
    n=np.array([[0,0,1]])
    invk=np.linalg.inv(K)
    le=np.expand_dims(K@r@invk,axis=0)
    inv_depths=np.expand_dims(inv_depths,axis=(1,2))
    homographies=K@((t@n)*inv_depths)@invk + le

    """ YOUR CODE ENDS HERE """

    return homographies

# Part 2
def compute_plane_sweep_volume(images, ref_pose, K, inv_depths, img_hw):
    """Compute plane sweep volume, by warping all images to the reference camera
    fronto-parallel planes, before computing the variance for each pixel and
    depth.

    Args:
        images (list[Image]): List of images which contains information about
          the camera extrinsics for each image
        ref_pose (np.ndarray): Reference camera pose
        K (np.ndarray): 3x3 intrinsic matrix (assumed same for all cameras)
        inv_depths (list): List of inverse depths to consider for plane sweep
        img_hw (tuple): tuple containing (H, W), which are the output height
          and width for the plane sweep volume.

    Returns:
        ps_volume (np.ndarray):
          Plane sweep volume of size (D, H, W), with dtype=np.float64, where
          D is len(inv_depths), and (H, W) are the image heights and width
          respectively. Each element should contain the variance of all pixel
          intensities that warp onto it.
        accum_count (np.ndarray):
          Accumulator count of same size as ps_volume, and dtype=np.int32.
          Keeps track of how many images are warped into a certain pixel,
          i.e. the number of pixels used to compute the variance.
        extras (any type):
          Any additional information you might want to keep for part 4.
    """
    """ YOUR CODE STARTS HERE """
    D = len(inv_depths)
    H, W = img_hw
    N=len(images)
    ps_volume = np.zeros((D, H, W), dtype=np.float32)
    ps_all = np.zeros([N,H,W,3], dtype=np.float32)
    accum_count = np.zeros((D, H, W), dtype=np.uint8)
    mask=np.zeros([N,H,W,3],dtype=np.uint8)
    extras = []

    H = np.zeros([N,D,3,3])

    def wvar(a, w):
        av = np.average(a, weights=w, axis=0)
        av=np.expand_dims(av,axis=0)
        var = np.average((a - av) ** 2, weights=w,axis=0)
        return var

    for i in range(N):
        relative_pose = compute_relative_pose(images[i].pose_mat, ref_pose)
        H[i] = get_plane_sweep_homographies(K, relative_pose, np.array(inv_depths))
    extras.append(H)
    boraderValue=-1
    for j in range(0,D):
        for i in range(N):
            ps_all[i]=cv2.warpPerspective(images[i].image,np.linalg.inv(H[i][j]),img_hw[::-1],
                                          borderValue=(boraderValue,boraderValue,boraderValue))
        mask[:, :, :, :] = (ps_all != boraderValue)
        accum_count[j]=np.sum(mask[:,:,:,0],axis=0)
        ps_volume[j]=np.mean(wvar(ps_all,w=mask),axis=2)
    """ YOUR CODE ENDS HERE """

    return ps_volume, accum_count, extras

def compute_depths(ps_volume, inv_depths):
    """Computes inverse depth map from plane sweep volume as the
    argmin over plane sweep volume variances.

    Args:
        ps_volume (np.ndarray): Plane sweep volume of size (D, H, W) from
          compute_plane_sweep_volume()
        inv_depths (np.ndarray): List of depths considered in the plane
          sweeping (D,)

    Returns:
        inv_depth_image (np.ndarray): inverse-depth estimate (H, W)
    """
    """ YOUR CODE STARTS HERE """
    idx=np.argmin(ps_volume,axis=0)
    inv_depth_image= inv_depths[idx]

    """ YOUR CODE ENDS HERE """

    return inv_depth_image

# Part 3
def unproject_depth_map(image, inv_depth_image, K, mask=None):
    """Converts the depth map into points by unprojecting depth map into 3D

    Note: You will also need to implement the case where no mask is provided

    Args:
        image (np.ndarray): Image bitmap (H, W, 3)
        inv_depth_image (np.ndarray): Inverse depth image (H, W)
        K (np.ndarray): 3x3 Camera intrinsics
        mask (np.ndarray): Optional mask of size (H, W) and dtype=np.bool.

    Returns:
        xyz (np.ndarray): Nx3 coordinates of points, dtype=np.float64.
        rgb (np.ndarray): Nx3 RGB colors, where rgb[i, :] is the (Red,Green,Blue)
          colors for the points at position xyz[i, :]. Should be in the range
          [0, 1] and have dtype=np.float64.
    """

    """ YOUR CODE STARTS HERE """
    H, W = image.shape[:-1]
    xx, yy = np.meshgrid(np.arange(W), np.arange(H))
    homocoor = np.concatenate([xx.reshape(-1, 1), yy.reshape(-1, 1), np.ones([H * W, 1])], axis=1)

    xyz = np.zeros([H * W, 3], dtype=np.float64)
    rgb = np.zeros([H * W, 3], dtype=np.float64)

    depth = inv_depth_image[yy.reshape(-1), xx.reshape(-1)]
    xyz[:, -1] = -depth / (homocoor @ np.linalg.inv(K).T)[:, -1]
    xyz[:, :-1] = (homocoor @ np.linalg.inv(K).T)[:, :-1]

    rgb[:, :] = image[yy.reshape(-1), xx.reshape(-1)]

    if (mask is not None):
        coormask = mask[yy.reshape(-1), xx.reshape(-1)]
        coormask = coormask.astype(np.bool)
        xyz = xyz[coormask]
        rgb = rgb[coormask]
    """ YOUR CODE ENDS HERE """

    return xyz, rgb

# Part 4
def post_process(ps_volume, inv_depth, accum_count, extras):
    """Post processes the plane sweep volume and compute a mask to indicate
    which pixels have confident estimates of the depth

    Args:
        ps_volume: Plane sweep volume from compute_plane_sweep_volume()
          of size (D, H, W)
        inv_depths (List[float]): List of depths considered in the plane
          sweeping
        accum_count: Accumulator count from compute_plane_sweep_volume(), which
          can be used to indicate which pixels are not observed by many other
          images.
        extras: Extra variables from compute_plane_sweep_volume() in Part 2

    Returns:
        inv_depth_image: Denoised Inverse depth image (similar to compute_depths)
        mask: np.ndarray of size (H, W) and dtype np.bool.
          Pixels with values TRUE indicate valid pixels.
    """

    """ YOUR CODE STARTS HERE """
    D, H, W = ps_volume.shape

    idx = np.argmin(ps_volume, axis=0)
    inv_depth_image = inv_depth[idx]
    accum_count = accum_count.transpose(2, 0, 1).reshape(-1, D)
    maskaccum = accum_count[np.arange(H * W), idx.reshape(-1)].reshape(H, W)

    idf = (inv_depth_image.copy() / 1.25 * 255).astype(np.uint8) / 255

    idf = cv2.fastNlMeansDenoising(
        (idf * 255).astype(np.uint8), None, h=35, templateWindowSize=9, searchWindowSize=21) / 255
    idf = cv2.medianBlur((idf * 255).astype(np.uint8), 9) / 255
    idf = cv2.GaussianBlur((idf * 255).astype(np.uint8), (9, 9), 0) / 255
    denoisedepth = idf.copy()
    confidentdepth = cv2.medianBlur((idf * 255).astype(np.uint8), 31) / 255
    dif = np.abs(inv_depth_image - denoisedepth) / confidentdepth
    dif[maskaccum <= 5] = 0
    median = np.median(dif)
    dif[dif >= median] = 1
    dif[dif < median] = 0

    mask = dif.astype(np.bool)
    """ YOUR CODE ENDS HERE """

    return idf*1.25, mask

reason: import numpy,import scipy,from scipy


text:
#
# For licensing see accompanying LICENSE file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#

import random
import torch
import numpy as np
import os
from typing import Union, Dict, Optional, Tuple
from torch import Tensor
from sys import platform

from utils import logger
from utils.ddp_utils import is_master
from cvnets.layers import norm_layers_tuple

def check_compatibility():
    ver = torch.__version__.split(".")
    major_version = int(ver[0])
    minor_version = int(ver[0])

    if major_version < 1 and minor_version < 7:
        logger.error(
            "Min pytorch version required is 1.7.0. Got: {}".format(".".join(ver))
        )

def check_frozen_norm_layer(model: torch.nn.Module) -> (bool, int):

    if hasattr(model, "module"):
        model = model.module

    count_norm = 0
    frozen_state = False
    for m in model.modules():
        if isinstance(m, norm_layers_tuple):
            frozen_state = m.weight.requires_grad

    return frozen_state, count_norm

def device_setup(opts):
    random_seed = getattr(opts, "common.seed", 0)
    random.seed(random_seed)
    torch.manual_seed(random_seed)
    np.random.seed(random_seed)

    is_master_node = is_master(opts)
    if is_master_node:
        logger.log("Random seeds are set to {}".format(random_seed))
        logger.log("Using PyTorch version {}".format(torch.__version__))

    n_gpus = torch.cuda.device_count()
    if n_gpus == 0:
        if is_master_node:
            logger.warning("No GPUs available. Using CPU")
        device = torch.device("cpu")
        n_gpus = 0
    else:
        if is_master_node:
            logger.log("Available GPUs: {}".format(n_gpus))
        device = torch.device("cuda")

        if torch.backends.cudnn.is_available():
            import torch.backends.cudnn as cudnn

            torch.backends.cudnn.enabled = True
            cudnn.benchmark = False
            cudnn.deterministic = True
            if is_master_node:
                logger.log("CUDNN is enabled")

    setattr(opts, "dev.device", device)
    setattr(opts, "dev.num_gpus", n_gpus)

    return opts

def create_directories(dir_path: str, is_master_node: bool) -> None:
    if not os.path.isdir(dir_path):
        os.makedirs(dir_path)
        if is_master_node:
            logger.log("Directory created at: {}".format(dir_path))
    else:
        if is_master_node:
            logger.log("Directory exists at: {}".format(dir_path))

def move_to_device(
    opts,
    x: Union[Dict, Tensor],
    device: Optional[str] = "cpu",
    non_blocking: Optional[bool] = True,
    *args,
    **kwargs
) -> Union[Dict, Tensor]:

    # if getattr(opts, "dataset.decode_data_on_gpu", False):
    #    # data is already on GPU
    #    return x

    if isinstance(x, Dict):
        # return the tensor because if its already on device
        if "on_gpu" in x and x["on_gpu"]:
            return x

        for k, v in x.items():
            if isinstance(v, Dict):
                x[k] = move_to_device(opts=opts, x=v, device=device)
            elif isinstance(v, Tensor):
                x[k] = v.to(device=device, non_blocking=non_blocking)

    elif isinstance(x, Tensor):
        x = x.to(device=device, non_blocking=non_blocking)
    else:
        logger.error(
            "Inputs of type  Tensor or Dict of Tensors are only supported right now"
        )
    return x

def is_coreml_conversion(opts) -> bool:
    coreml_convert = getattr(opts, "common.enable_coreml_compatible_module", False)
    if coreml_convert:
        return True
    return False

ext: py

max_stars_repo_name: apple/ml-cvnets

max_issues_repo_name: apple/ml-cvnets

max_forks_repo_name: apple/ml-cvnets

reason: import numpy