|
|
|
|
|
import json |
|
import requests |
|
import shutil |
|
from pathlib import Path |
|
|
|
import cv2 |
|
import numpy as np |
|
import torch |
|
from tqdm.auto import tqdm |
|
|
|
import logger |
|
|
|
DATA_URL = "https://cvg-data.inf.ethz.ch/OrienterNet_CVPR2023" |
|
|
|
|
|
def read_image(path, grayscale=False): |
|
if grayscale: |
|
mode = cv2.IMREAD_GRAYSCALE |
|
else: |
|
mode = cv2.IMREAD_COLOR |
|
image = cv2.imread(str(path), mode) |
|
if image is None: |
|
raise ValueError(f"Cannot read image {path}.") |
|
if not grayscale and len(image.shape) == 3: |
|
image = np.ascontiguousarray(image[:, :, ::-1]) |
|
return image |
|
|
|
|
|
def write_torch_image(path, image): |
|
image_cv2 = np.round(image.clip(0, 1) * 255).astype(int)[..., ::-1] |
|
cv2.imwrite(str(path), image_cv2) |
|
|
|
|
|
class JSONEncoder(json.JSONEncoder): |
|
def default(self, obj): |
|
if isinstance(obj, (np.ndarray, torch.Tensor)): |
|
return obj.tolist() |
|
elif isinstance(obj, np.generic): |
|
return obj.item() |
|
return json.JSONEncoder.default(self, obj) |
|
|
|
|
|
def write_json(path, data): |
|
with open(path, "w") as f: |
|
json.dump(data, f, cls=JSONEncoder) |
|
|
|
|
|
def download_file(url, path): |
|
path = Path(path) |
|
if path.is_dir(): |
|
path = path / Path(url).name |
|
path.parent.mkdir(exist_ok=True, parents=True) |
|
logger.info("Downloading %s to %s.", url, path) |
|
with requests.get(url, stream=True) as r: |
|
total_length = int(r.headers.get("Content-Length")) |
|
with tqdm.wrapattr(r.raw, "read", total=total_length, desc="") as raw: |
|
with open(path, "wb") as output: |
|
shutil.copyfileobj(raw, output) |
|
return path |
|
|
