import torch
import torch.nn as nn


class CameraModule(nn.Module):
    def __init__(self, image_size):
        super(CameraModule, self).__init__()
        self.image_size = image_size

    def perspective(self, query_pts, calibrations):
        query_pts = torch.bmm(calibrations[:, :3, :3], query_pts)
        query_pts = query_pts + calibrations[:, :3, 3:4]
        query_pts_xy = query_pts[:, :2, :] / query_pts[:, 2:, :]
        query_pts_xy = query_pts_xy
        return query_pts_xy