|
|
|
|
|
import json
|
|
import requests
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from tqdm.auto import tqdm
|
|
|
|
from .. 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
|
|
|
