import os
import sys
import re
import numpy as np
import cv2
import json
import yaml
import vis_utils as v_uts
import struct
from cv_base import (
Faces, Aux, Obj, DEFAULT_MATERIAL
)
hasTorch = True
try:
import torch
except:
hasTorch = False
import functools
import pandas as pd
from tqdm import tqdm
from PIL import Image
try:
from plyfile import PlyData
except:
"no ply"
import pdb
b=pdb.set_trace
def default(x, val):
return val if x is None else x
class IOShop:
def __init__(self, name, **kwargs):
ioFuncs = {'depth': DepthIO,
'image': ImageIO,
'flow': FlowIO,
'segment': SegmentIO,
'prob': ProbIO,
'video': VideoIO}
self.io = ioFuncs[name](**kwargs)
def load(self, file_name, **kwargs):
return self.io.load(file_name, **kwargs)
def dump(self, file_name, file, **kwargs):
self.io.dump(file_name, file, **kwargs)
class BaseIO:
def __init__(self, appex='jpg'):
self.type = 'image'
self.appex = appex
def load(self, file_name):
file_name = '%s.%s' % (file_name, self.appex)
image = cv2.imread(file_name, cv2.IMREAD_UNCHANGED)
assert not (image is None), '%s not exists' % file_name
return image
def dump(self, file_name, file):
v_uts.mkdir_if_need(os.path.dirname(file_name))
file_name = '%s.%s' % (file_name, self.appex)
cv2.imwrite(file_name, file)
class ImageIO(BaseIO):
def __init__(self, appex='jpg'):
super(ImageIO, self).__init__(appex=appex)
self.type = 'image'
def load(self, file_name):
if file_name.endswith('heic') or file_name.endswith('HEIC'):
byte = read2byte(file_name)
image = decodeImage(byte)
else:
image = super(ImageIO, self).load(file_name)
return image
@staticmethod
def imwrite(file_name, data, order='rgb'):
cv2.imwrite(file_name, data[:, :, ::-1])
class SegmentIO(BaseIO):
def __init__(self):
super(SegmentIO, self).__init__(appex='png')
self.type = 'segment'
class ProbIO(BaseIO):
def __init__(self):
super(ProbIO, self).__init__()
self.type = 'prob'
self.max_class = 4
def load(self, file_name, channels=None):
image = cv2.imread(file_name, cv2.IMREAD_UNCHANGED)
channels = default(channels, self.max_class)
output = np.zeros(image.shape[:2])
# for i in range(channels):
def dump(self, file_name, file):
"""
height, width, channel
"""
output = np.zeros((height, width), dtype=np.uint16)
h, w, c = file.shape
for i in range(c):
output = output + np.uint16(file[:, :, i] * 255) + i * 256
cv2.imwrite(file_name, output.astype('uint16'))
class MeshIO(BaseIO):
def __init__(self):
super().__init__(appex='obj')
self.type = 'mesh'
def dump_obj(self, filename, obj):
export_obj(filename, obj)
def load_obj(self, filename):
return load_obj(filename)
def normalize_normal(mat):
mat = (mat / 255.0 * 2.0 - 1.0).astype('float32')
l1 = np.linalg.norm(mat, axis=2)
for j in range(3):
mat[:,:,j] /= (l1 + 1e-9)
return mat
class NormalIO(BaseIO):
def __init__(self, xyz='rgb'):
"""
rgb: means the normal saved in the order of x: r ...
"""
self._xyz = xyz
def read(self, filename):
normal = cv2.imread(filename, cv2.IMREAD_UNCHANGED)
if self._xyz == 'rgb':
normal = normal[:, :, ::-1]
normal = normalize_normal(normal)
return normal
class DepthIO(BaseIO):
def __init__(self, bit=8):
super(DepthIO, self).__init__(appex='pfm')
assert bit in [8, 16]
scale = {8: 1, 16: 2}
self.bits = scale[bit]
self.dump_vis = True
def load(self, path):
"""Read pfm file.
Args:
path (str): path to file
Returns:
tuple: (data, scale)
"""
path = '%s.%s' % (path, self.appex)
with open(path, "rb") as file:
color = None
width = None
height = None
scale = None
endian = None
header = file.readline().rstrip()
if header.decode("ascii") == "PF":
color = True
elif header.decode("ascii") == "Pf":
color = False
else:
raise Exception("Not a PFM file: " + path)
dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
if dim_match:
width, height = list(map(int, dim_match.groups()))
else:
raise Exception("Malformed PFM header.")
scale = float(file.readline().decode("ascii").rstrip())
if scale < 0:
# little-endian
endian = "<"
scale = -scale
else:
# big-endian
endian = ">"
data = np.fromfile(file, endian + "f")
shape = (height, width, 3) if color else (height, width)
data = np.reshape(data, shape)
data = np.flipud(data)
return data, scale
def dump(self, path, image, scale=1):
"""Write pfm file.
Args:
path (str): pathto file
image (array): data
scale (int, optional): Scale. Defaults to 1.
"""
v_uts.mkdir_if_need(os.path.dirname(path))
path = path + '.pfm'
with open(path, "wb") as file:
color = None
if image.dtype.name != "float32":
raise Exception("Image dtype must be float32.")
image = np.flipud(image)
if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True
elif (
len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
): # greyscale
color = False
else:
raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
file.write("PF\n" if color else "Pf\n".encode())
file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
endian = image.dtype.byteorder
if endian == "<" or endian == "=" and sys.byteorder == "little":
scale = -scale
file.write("%f\n".encode() % scale)
image.tofile(file)
if self.dump_vis:
self.dump_visualize(path[:-4], image, self.bits)
@staticmethod
def to8UC3(depth, scale=1000):
"""
Convert depth image to 8UC3 format.
"""
h, w = depth.shape
max_depth = (256.0 ** 3 - 1) / scale
# Clip depth values exceeding the maximum depth
depth = np.clip(depth, 0, max_depth)
# Scale the depth values
value = depth * scale
# Split the depth values into three channels
ch = np.zeros((h, w, 3), dtype=np.uint8)
ch[:, :, 0] = np.uint8(value / (256 ** 2))
ch[:, :, 1] = np.uint8((value % (256 ** 2)) / 256)
ch[:, :, 2] = np.uint8(value % 256)
return ch
@staticmethod
def read8UC3(depth, scale=1000):
"""
Convert 8UC3 image to scaled depth representation.
"""
if isinstance(depth, str):
depth = cv2.imread(depth, cv2.IMREAD_UNCHANGED)
# Merge the three channels into a single depth value
depth_uint16 = depth[:, :, 0] * (256 ** 2) + \
depth[:, :, 1] * 256 + depth[:, :, 2]
# Convert depth to the scaled representation
depth = depth_uint16.astype(np.float32) / scale
return depth
@staticmethod
def dump_visualize(path, depth, bits=1):
depth_min = depth.min()
depth_max = depth.max()
max_val = (2**(8*bits))-1
if depth_max - depth_min > np.finfo("float").eps:
out = max_val * (depth - depth_min) / (depth_max - depth_min)
else:
out = 0
if bits == 1:
cv2.imwrite(path + ".png", out.astype("uint8"))
elif bits == 2:
cv2.imwrite(path + ".png", out.astype("uint16"))
return
@staticmethod
def load_png(path):
depth = cv2.imread(path, cv2.IMREAD_UNCHANGED)
return depth
@staticmethod
def dump_png(path, depth, bits=2, max_depth=20.0):
assert (path.endswith(".png"))
max_val = (2**(8*bits))-1
depth = depth / max_depth * max_val
cv2.imwrite(path, depth.astype("uint16"))
@staticmethod
def read_depth(filename, scale=6000, sz=None, is_disparity=False):
if not hasTorch:
return None
depth = cv2.imread(filename, cv2.IMREAD_UNCHANGED)
depth = np.float32(depth) / scale
if sz:
h, w = sz
depth = cv2.resize(depth, (w, h),
interpolation=cv2.INTER_NEAREST)
depth = torch.from_numpy(depth)
if is_disparity: # convert to depth
depth = 1.0 / torch.clamp(depth, min=1e-10)
return depth
def write_depth(path, depth, grayscale, bits=1):
"""Write depth map to png file.
Args:
path (str): filepath without extension
depth (array): depth
grayscale (bool): use a grayscale colormap?
"""
if not grayscale:
bits = 1
if not np.isfinite(depth).all():
depth=np.nan_to_num(depth, nan=0.0, posinf=0.0, neginf=0.0)
print("WARNING: Non-finite depth values present")
depth_min = depth.min()
depth_max = depth.max()
max_val = (2**(8*bits))-1
if depth_max - depth_min > np.finfo("float").eps:
out = max_val * (depth - depth_min) / (depth_max - depth_min)
else:
out = np.zeros(depth.shape, dtype=depth.dtype)
if not grayscale:
out = cv2.applyColorMap(np.uint8(out), cv2.COLORMAP_INFERNO)
if bits == 1:
cv2.imwrite(path + ".png", out.astype("uint8"))
elif bits == 2:
cv2.imwrite(path + ".png", out.astype("uint16"))
return
class NormalIO(BaseIO):
def __init__(self):
super(NormalIO, self).__init__(appex='npy')
self.dump_vis = False
@staticmethod
def read_normal(filename, sz=None, to_torch=False):
if not hasTorch:
return None
if not os.path.exists(filename):
h, w = sz
return torch.ones((h, w, 3)) * 0.3
image = cv2.imread(filename)[:, :, ::-1]
image = np.float32(image)
image = (image / 127.5 - 1)
if sz:
h, w = sz
image = cv2.resize(image, (w, h),
interpolation=cv2.INTER_NEAREST)
return torch.from_numpy(image)
def to8UC3(self, normal):
return np.uint8((normal + 1) * 127.5)
class FlowIO(BaseIO):
def __init__(self):
super(FlowIO, self).__init__(appex='npy')
self.dump_vis = False
def normalize(self, flow, shape=None):
if shape is None:
shape = flow.shape[:2]
flow[:, :, 0] /= shape[1]
flow[:, :, 1] /= shape[0]
return flow
def denormalize(self, flow, shape=None):
if shape is None:
shape = flow.shape[:2]
flow[:, :, 0] *= shape[1]
flow[:, :, 1] *= shape[0]
return flow
def visualization(self, flow):
pass
def load(self, path, shape=None):
path = path + '.npy'
flow = np.load(path)
flow = self.denormalize(flow, shape)
assert flow is not None
return flow
def dump(self, path, flow):
v_uts.mkdir_if_need(os.path.dirname(path))
path = path + '.npy'
flow = self.normalize(flow)
np.save(path, flow)
if self.dump_vis:
self.dump_visualize(path[:-4], flow)
def dump_visualize(self, path, flow):
_, flow_c = v_uts.flow2color(flow)
cv2.imwrite(path + '.png', flow_c)
class VideoIO(BaseIO):
def __init__(self, longside_len=None):
super(VideoIO, self).__init__()
self.longside_len = longside_len
def get_fps(self, path):
vidcap = cv2.VideoCapture(path)
return vidcap.get(cv2.CAP_PROP_FPS)
def load_first_frame(self, path):
import skvideo.io as vio
video = vio.vreader(path)
frame = next(video)
if self.longside_len is not None:
frame = v_uts.resize2maxsize(frame, self.longside_len)
return frame
def load(self, path, sample_rate=1, max_len=1e10,
load_to_dir=False,
dir_name=None,
pre_len=5,
save_transform=None):
import skvideo.io as vio
def default_transform(x):
if x.ndim == 2:
return x
if x.ndim == 3 and x.shape[2] == 3:
return x[:, :, ::-1]
return x
frames = []
reader = vio.vreader(path)
if load_to_dir:
v_uts.mkdir(dir_name)
if save_transform is None:
save_transform = lambda x : x
for count, frame in enumerate(reader):
if count == max_len:
break
if count % sample_rate == 0:
if self.longside_len is not None:
frame = v_uts.resize2maxsize(
frame, self.longside_len)
if load_to_dir:
img_file = f"{dir_name}/{count:05}.png"
frame = save_transform(frame)
cv2.imwrite(img_file, frame)
else:
frames.append(frame)
if not load_to_dir:
return frames
def load_till_end(self, path, sample_rate=1):
import skvideo.io as vio
frames = []
reader = vio.vreader(path)
count = 0
while True:
try:
frame = next(reader)
except:
break
if count % sample_rate == 0:
if self.longside_len is not None:
frame = v_uts.resize2maxsize(
frame, self.longside_len)
frames.append(frame)
count += 1
return frames
def load_w_cv(self, path, out_dir, sample_rate = 1, ext="jpg"):
v_uts.video_to_frame(path,
out_dir,
max_len=self.longside_len,
sample_rate=sample_rate,
ext=ext)
def dump_to_images(self, frames, image_path):
v_uts.mkdir_if_need(image_path)
for count, frame in tqdm(enumerate(frames)):
image_file = '%s/%04d.jpg' % (image_path, count)
cv2.imwrite(image_file, frame[:, :, ::-1])
def dump(self, path, frames, fps=30, lossless=False):
from moviepy.editor import ImageSequenceClip, VideoFileClip
if isinstance(frames[0], str):
frame_np = []
for frame in tqdm(frames):
cur_frame = cv2.imread(frame, cv2.IMREAD_UNCHANGED)[:, :, ::-1]
frame_np.append(cur_frame)
frames = frame_np
clip = ImageSequenceClip(frames, fps)
if lossless:
assert path.endswith('avi')
clip.write_videofile(path, codec='png')
else:
clip.write_videofile(path, fps=fps)
def dump_skv(self, path, frames, fps=30):
if frames[0].ndim == 2:
frames = [cv2.cvtColor(frame,cv2.COLOR_GRAY2RGB) for frame in frames]
else:
frames = [frame[:, :, ::-1] for frame in frames]
v_uts.frame_to_video_simple(frames, fps, video_name=path)
# import skvideo.io as vio
# fps = str(int(fps))
# vid_out = vio.FFmpegWriter(path,
# inputdict={'-r': fps},
# outputdict={
# '-vcodec': 'libx264',
# '-pix_fmt': 'yuv420p',
# '-r': fps,
# },
# verbosity=1)
# for idx, frame in enumerate(frames):
# vid_out.writeFrame(frame)
# vid_out.close()
def resave_video(self, video_file, start, end,
outvideo_file):
"""
:param start: sec start
:param end: sec end
:return:
"""
fps = self.get_fps(video_file)
frames = self.load(video_file)
start_frame = int(start * fps)
end_frame = int(end * fps)
frames = frames[start_frame:end_frame]
self.dump_skv(outvideo_file, frames, fps)
def frame2video(self, folder, output, ext=".jpg"):
image_files = v_uts.list_all_files(folder, exts=[ext])
frames = []
for name in tqdm(image_files):
frames.append(cv2.imread(name)[:, :, ::-1])
self.dump(output, frames)
class NpEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, np.integer):
return int(obj)
if isinstance(obj, np.floating):
return float(obj)
if isinstance(obj, np.ndarray):
return obj.tolist()
return super(NpEncoder, self).default(obj)
def read2byte(filename):
with open(filename, 'rb') as f:
file_data = f.read()
return file_data
def decodeImage(bytesIo):
import whatimage
import pyheif
from PIL import Image
fmt = whatimage.identify_image(bytesIo)
if fmt in ['heic', 'avif']:
i = pyheif.read_heif(bytesIo)
# Convert to other file format like jpeg
pi = Image.frombytes(
mode=i.mode, size=i.size, data=i.data)
image = np.asarray(pi)
image = image[:, :, ::-1] # to BGR
return image
else:
return None
def image2Normal(imagePath):
from skimage import io
normal = io.imread(imagePath)
normal = ((np.float32(normal) / 255.0) * 2 - 1.0 )
return normal
def normal2Image(normal):
nm_pred_val = (normal + 1.) / 2.
nm_pred_val = np.uint8(nm_pred_val*255.)
return nm_pred_val
def dump_normal(filename, normal):
normal = normal2Image(normal)
cv2.imwrite(filename + '.png', array)
def dump_prob2image(filename, array):
"""
dump probility map to image when
array: [x, height, width] (x = 1, 3, 4)
"""
class_num = array.shape[0]
# assert class_num <= 4
if class_num >= 4 :
print('warning: only save the first 3 channels')
array = array[:3, :, :]
if class_num == 2:
raise ValueError('not implement')
array = np.transpose(np.uint8(array * 255), (1, 2, 0))
if filename.endswith('.png'):
cv2.imwrite(filename, array)
return
cv2.imwrite(filename + '.png', array)
assert os.path.exists(filename)
def load_image2prob(filename):
if not filename.endswith('.png'):
filename = filename + '.png'
array = cv2.imread(filename, cv2.IMREAD_UNCHANGED)
array = np.transpose(array, (2, 0, 1)) / 255
return array
def shape_match(images):
assert len(images) > 1
shape = images[0].shape[:2]
for image in images[1:]:
cur_shape = image.shape[:2]
if np.sum(np.abs(np.array(shape) - \
np.array(cur_shape))):
return False
return True
def append_apex(filename, appex):
filename = filename.split('.')
prefix = '.'.join(filename[:-1])
filetype = filename[-1]
return '%s_%s.%s' % (prefix, appex, filetype)
def load_json(json_file):
with open(json_file) as f:
res = json.load(f)
return res
def dump_numpy(filename, x: np.ndarray):
np.savetxt(filename, x, delimiter=' ', fmt='%1.6f')
def dump_json(filename, odgt, w_np=False):
with open(filename, 'w') as f:
if not w_np:
json.dump(odgt, f, indent=4)
else:
json.dump(odgt, f, indent=4, cls=NpEncoder)
def dump_jsonl(filename, odgt):
with open(filename, 'w') as file:
for entry in odgt:
json.dump(entry, file)
file.write('\n')
def dump_pair_data(image_list,
label_list,
outfile,
root='',
data_type='txt',
fields=None):
if fields is None:
fields = ["image", "segment"]
if data_type == 'txt':
fp = open(outfile, 'w')
for imagefile, labelfile in zip(image_list, label_list):
imagefile = imagefile.replace(root, '.')
labelfile = labelfile.replace(root, '.')
fp.write('%s %s\n' % (imagefile, labelfile))
fp.close()
elif data_type == "odgt":
odgt = []
for imagefile, labelfile in zip(image_list, label_list):
imagefile = imagefile.replace(root, '.')
labelfile = labelfile.replace(root, '.')
item = {fields[0]: imagefile,
fields[1]: labelfile}
odgt.append(item)
dump_json(outfile, odgt)
def save_xlsx(filename, dicts, sheets=None):
"""
Save a list of dicts to an xlsx file.
"""
with pd.ExcelWriter(filename, mode='w') as writer:
if sheets is None:
df1 = pd.DataFrame(dicts)
df1.to_excel(writer, index=False)
return
for sheet in sheets:
df1 = pd.DataFrame(dicts[sheet])
df1.to_excel(writer, sheet_name=sheet, index=False)
def load_xlsx(filename, sheets=None):
assert os.path.exists(filename) , f"File not found: {filename}"
if sheets is None:
df = pd.read_excel(filename)
dict = {}
for column in df.columns:
dict[column] = df[column].tolist()
else:
dict = {}
for sheet in sheets:
df = pd.read_excel(filename, sheet_name=sheet)
cur_dict = {}
for column in df.columns:
cur_dict[column] = df[column].tolist()
print(cur_dict.keys())
dict[sheet] = cur_dict
print(dict.keys())
return dict
def dump_lines(filename, file_list):
f = open(filename, 'w')
tbar = tqdm(file_list)
for i, elements in enumerate(tbar):
if isinstance(elements, (tuple, list)):
line = ' '.join(elements)
elif isinstance(elements, str):
line = elements
appex = '' if i == len(file_list) - 1 else '\n'
f.write('%s%s' % (line, appex))
f.close()
def load_lines(txt_file):
lines = [line.strip() for line in open(txt_file, 'r')]
return lines
def load_jsonl(jsonl_file):
# List to hold all JSON objects
data = []
# Open the file and read line by line
with open(jsonl_file, 'r') as file:
for line in file:
# Each line is a JSON object, parse it and append to the list
json_object = json.loads(line)
data.append(json_object)
return data
def load_yaml(yaml_file):
with open(yaml_file, "r") as f:
yaml_dict = yaml.safe_load(f)
return yaml_dict
def load_odgt(odgt):
try:
samples = [json.loads(x.rstrip()) \
for x in open(odgt, 'r')][0]
except:
samples = load_json(odgt)
print(samples[0].keys())
return samples
def fuse_odgt(odgt_files):
"""
odgt_files:
"""
odgt_full = []
for odgt_file in odgt_files:
odgt = load_odgt(odgt_file)
odgt_full = odgt_full + odgt
return odgt_full
def load_video_first_frame(video_name):
cap = cv2.VideoCapture(video_name)
if(cap.isOpened()):
ret, frame = cap.read()
else:
raise ValueError("can not read %s" % video_name)
return frame
def load_lines(txt_file):
lines = [line.strip() for line in open(txt_file, 'r')]
return lines
def load_csv(csv_file):
import csv
lines = []
with open(csv_file) as f:
reader = csv.reader(f, delimiter=',')
for row in reader:
lines.append(row)
return lines[1:]
# cat multi files in to a single file
def cat_files(files, output):
all_lines = []
for filename in files:
lines = load_lines(filename)
all_lines = all_lines + lines
dump_lines(output, all_lines)
class SkipExist:
def __init__(self,
processor,
ioType='image',
need_res=False,
rerun=False):
self.ioType = ioType
self.io = IOShop(self.ioType).io
self.processor = processor
self.rerun = rerun
self.need_res = need_res
def __call__(self, *args, **kwargs):
assert 'filename' in kwargs
true_file = '%s.%s' % (kwargs['filename'], self.io.appex)
if os.path.exists(true_file):
if self.need_res:
res = self.io.load(kwargs['filename'])
return res
else:
filename = kwargs['filename']
del kwargs['filename']
res = self.processor(*args, **kwargs)
self.io.dump(filename, res)
def dump_pkl(filename, data):
import pickle as pkl
with open(filename, "wb") as fl:
pkl.dump(data, fl)
def load_pkl(filename):
import pickle as pkl
with open(filename, 'rb') as fl:
res = pkl.load(fl)
return res
def write_pointcloud(filename, xyz_points, faces=None, rgb_points=None):
"""
creates a .pkl file of the point clouds generated
"""
assert xyz_points.shape[1] == 3,'Input XYZ points should be Nx3 float array'
if rgb_points is None:
rgb_points = np.ones(xyz_points.shape).astype(np.uint8) * 255
else:
rgb_points = rgb_points.astype(np.uint8)
assert xyz_points.shape == rgb_points.shape,\
f'Input RGB colors should be Nx3 {rgb_points.shape} float array \
and have same size as input XYZ points {xyz_points.shape}'
# Write header of .ply file
fid = open(filename,'wb')
fid.write(bytes('ply\n', 'utf-8'))
fid.write(bytes('format binary_little_endian 1.0\n', 'utf-8'))
fid.write(bytes('element vertex %d\n'%xyz_points.shape[0], 'utf-8'))
fid.write(bytes('property float x\n', 'utf-8'))
fid.write(bytes('property float y\n', 'utf-8'))
fid.write(bytes('property float z\n', 'utf-8'))
fid.write(bytes('property uchar red\n', 'utf-8'))
fid.write(bytes('property uchar green\n', 'utf-8'))
fid.write(bytes('property uchar blue\n', 'utf-8'))
fid.write(bytes('end_header\n', 'utf-8'))
# Write 3D points to .ply file
for i in range(xyz_points.shape[0]):
fid.write(bytearray(struct.pack("fffccc",xyz_points[i,0],xyz_points[i,1],xyz_points[i,2],
rgb_points[i,0].tostring(),rgb_points[i,1].tostring(),
rgb_points[i,2].tostring())))
if faces is not None:
for face in faces:
fid.write(struct.pack("<B", face[0]))
fid.write(struct.pack("<{}i".format(face[0]), *face[1]))
fid.close()
def read_ply(filename):
# Load the PLY file
ply_data = PlyData.read(filename)
# Access the vertex data
vertex_data = ply_data['vertex']
# Extract x, y, z coordinates as a numpy array
points = np.vstack((vertex_data['x'], vertex_data['y'], vertex_data['z'])).T
return points
def load_obj(file_path):
verts = []
normals = []
uvs = []
material_colors = []
texture_images = []
texture_atlas = []
faces_verts = []
faces_normals = []
faces_textures = []
faces_materials = []
with open(file_path, 'r') as file:
for line in file:
if line.startswith('v '):
vertex = [float(v) for v in line.split()[1:]]
verts.append(vertex)
elif line.startswith('vn '):
normal = [float(n) for n in line.split()[1:]]
normals.append(normal)
elif line.startswith('vt '):
uv = [float(u) for u in line.split()[1:]]
uvs.append(uv)
elif line.startswith("mtllib "):
mtl_name = line.split()[1]
elif line.startswith('vc '):
color = [float(c) for c in line.split()[1:]]
material_colors.append(color)
elif line.startswith('usemtl '):
material = line.split()[1]
texture_images.append(material)
elif line.startswith('f '):
face_data = line.split()[1:]
face_verts = []
face_normals = []
face_textures = []
for face in face_data:
res = face.split('/')
vert = res[0]
face_verts.append(int(vert))
if len(res) == 2:
texture = res[1]
face_textures.append(int(texture))
if len(res) == 3:
normal = res[2]
face_normals.append(int(normal))
faces_verts.append(face_verts)
faces_normals.append(face_normals)
faces_textures.append(face_textures)
faces_materials.append(len(texture_images) - 1)
mtl_file = f"{os.path.dirname(file_path)}/{mtl_name}"
with open(mtl_file, 'r') as file:
for line in file:
if line.startswith("map_Kd"):
image_name = line.split()[1]
break
assert len(texture_images) == 1
texture_name = texture_images[0]
image = cv2.imread(f"{os.path.dirname(file_path)}/{image_name}")
properties = Aux(
normals=np.array(normals),
verts_uvs=np.array(uvs),
material_colors=DEFAULT_MATERIAL,
texture_images={texture_name: np.float32(image)/ 255.0},
texture_atlas=None)
faces_verts=np.array(faces_verts)
num_faces = faces_verts.shape[0]
faces = Faces(
verts_idx=faces_verts,
normals_idx=np.ones(faces_verts.shape) * -1,
textures_idx=np.array(faces_textures),
materials_idx=np.zeros(num_faces))
obj = Obj(np.array(verts), faces, properties)
return obj
def export_obj(filename, obj,
include_normals=False,
include_textures=True):
"""
Export the given object to an .obj file with optional normals and textures.
Args:
filename (str): Path to the output .obj file (without the extension).
obj (namedtuple): Object containing vertices, faces, and properties.
include_normals (bool): Flag to include normals in the .obj file.
include_textures (bool): Flag to include textures in the .obj file.
"""
material_name = list(obj.properties.texture_images.keys())[0]
# Write obj file
name = os.path.basename(filename)
with open(filename + ".obj", "w") as f:
f.write("\n")
if include_textures:
f.write(f"mtllib {name}.mtl\n")
f.write("\n")
for vert in obj.verts:
x, y, z = vert
f.write(f"v {x} {y} {z}\n")
if include_textures:
for uv in obj.properties.verts_uvs:
x, y = uv
f.write(f"vt {x} {y}\n")
f.write(f"usemtl {material_name}\n")
num_faces = obj.faces.verts_idx.shape[0]
for i in range(num_faces):
f0, f1, f2 = obj.faces.verts_idx[i]
if include_textures:
t0, t1, t2 = obj.faces.textures_idx[i]
if t0 == -1:
f.write(f"f {f0} {f1} {f2}\n")
continue
f.write(f"f {f0}/{t0} {f1}/{t1} {f2}/{t2}\n")
else:
f.write(f"f {f0} {f1} {f2}\n")
# Write mtl file
if include_textures:
output_dir = os.path.dirname(filename)
with open(f"{output_dir}/{name}.mtl", "w") as f:
f.write(f"newmtl {material_name}\n")
f.write(f"map_Kd {name}.png\n")
material_colors = obj.properties.material_colors[material_name]
r, g, b = material_colors["ambient_color"]
f.write(f"Ka {r} {g} {b}\n")
r, g, b = material_colors["diffuse_color"]
f.write(f"Kd {r} {g} {b}\n")
r, g, b = material_colors["specular_color"]
f.write(f"Ks {r} {g} {b}\n")
s = material_colors["shininess"]
f.write(f"Ns {s}\n")
# Save texture image
image = obj.properties.texture_images[material_name] * 255
texture_img = f"{output_dir}/{name}.png"
cv2.imwrite(texture_img, image)
return
def resave_to_video():
folder = "/Users/peng/Downloads/DenseAR/Mesh/"
vname = "0037438511"
image_num = 125
frames = []
d_frames = []
crop = [0, 650, 1080, 1270]
for i in tqdm(range(image_num)):
name = f"{folder}/{vname}/{i}.jpg"
d_name = f"{folder}/{vname}/{i}.tiff"
img = np.array(Image.open(name))
depth = np.array(Image.open(d_name))
if img is None:
continue
img = img[crop[0]:crop[2], crop[1]:crop[3]]
depth = depth[crop[0]:crop[2], crop[1]:crop[3]]
depth = 1.0 / np.maximum(depth, 1e-10)
depth = p_uts.depth2color(depth, max_d=50)
frames.append(img)
d_frames.append(depth)
vio = io_uts.VideoIO()
video_file = f"{folder}/{vname}.mp4"
d_video_file = f"{folder}/{vname}_d.mp4"
vio.dump_skv(video_file, frames, fps=24)
vio.dump_skv(d_video_file, d_frames, fps=24)
def test_depth_8uc3_encode():
depth = np.random.rand(480, 640) * 200
dio = DepthIO()
depth_encode = dio.to8UC3(depth)
depth_decode = dio.read8UC3(depth_encode)
print(depth, depth_decode)
assert np.sum(np.abs(depth - depth_decode)) / (480 * 640) < 1e-3
########### copy from gta code ################
@functools.lru_cache()
def build_mesh(w, h):
w = np.linspace(-1.0, 1.0, num=w, dtype=np.float32)
h = np.linspace(1.0, -1.0, num=h, dtype=np.float32)
return np.stack(np.meshgrid(w, h), axis=0)
def build_proj_matrix(fov, aspect):
proj = np.zeros((4, 4))
proj[0, 0] = 1.0 / np.tan(np.radians(fov / 2)) / aspect
proj[1, 1] = 1.0 / np.tan(np.radians(fov / 2))
proj[2, 2] = 0.00001502 # reverse-engineered get from shader
proj[2, 3] = 0.15000225 # reverse-engineered get from shader
proj[3, 2] = -1.0
return proj
def zbuffer_to_depth(zbuffer, fov):
height, width = zbuffer.shape[:2]
aspect = width / height
mesh = build_mesh(width, height)
if len(zbuffer.shape) != 3:
zbuffer = np.expand_dims(zbuffer, 0)
pcloud = np.concatenate((mesh, zbuffer, np.ones_like(zbuffer)), 0)
pcloud = pcloud.reshape(4, height * width)
proj_matrix = build_proj_matrix(fov, aspect)
pcloud = np.linalg.inv(proj_matrix) @ pcloud
depth = -pcloud[2] / pcloud[3]
focal_cv = proj_matrix[0, 0] * width / 2.0
return depth.reshape(height, width), focal_cv
def test_zbuffer_to_depth():
# root = "E:/Dataset/GTA/Stereo_0/"
# name = root + "1-130423915874"
name = "E:/depth_video/0036696165/1"
config = load_json(name + ".json")
fov = config["fov"]
zbuffer = cv2.imread(name + ".tiff", cv2.IMREAD_UNCHANGED)
depth, focal = zbuffer_to_depth(zbuffer, fov)
print(depth)
def fuse_frames_of_depth_video():
"""
frames: list of images or video
"""
def frame_to_video(video_dir, video_name):
frames = v_uts.list_all_files(video_dir, exts=['jpg'])
rgb_video = f"{video_name}.mp4"
depth_video = f"{video_name}_d.avi"
cam_file = f"{video_name}.json"
dio = DepthIO()
imgs = []
depths = []
cams = []
print("seq len:", len(frames))
for i, frame in tqdm(enumerate(frames)):
name = f"{video_dir}/{i}.jpg"
d_name = f"{video_dir}/{i}.tiff"
c_name = f"{video_dir}/{i}.json"
img = np.array(Image.open(name))
depth = np.array(Image.open(d_name))
cam = load_json(c_name)
depth, focal = zbuffer_to_depth(depth, cam['fov'])
depth = dio.to8UC3(depth)
imgs.append(img)
depths.append(depth)
cam['focal'] = focal
cams.append(cam)
# if i > 30:
# break
vio = VideoIO()
vio.dump(rgb_video, imgs)
vio.dump(depth_video, depths, lossless=True)
dump_json(cam_file, cams)
folder = "E:/depth_video/"
output = "E:/depth_video_resave/"
v_uts.mkdir_if_need(output)
folder_names = v_uts.list_all_folders(folder)
for folder_name in tqdm(folder_names[1:]):
folder_name = folder_name.replace('\\', '/')
vid_name = folder_name.split('/')[-2]
print(folder_name, vid_name)
output_video = f"{output}/{vid_name}"
frame_to_video(folder_name, video_name=output_video)
# break
def save_xlsx(filename, dicts, sheets=None):
with pd.ExcelWriter(filename, mode='w') as writer:
if sheets is None:
df1 = pd.DataFrame(dicts)
df1.to_excel(writer, index=False)
return
for sheet in sheets:
df1 = pd.DataFrame(dicts[sheet])
df1.to_excel(writer, sheet_name=sheet, index=False)
def load_xlsx(filename, sheets=None):
assert os.path.exists(filename) , f"File not found: {filename}"
if sheets is None:
df = pd.read_excel(filename)
dict = {}
for column in df.columns:
dict[column] = df[column].tolist()
else:
dict = {}
for sheet in sheets:
df = pd.read_excel(filename, sheet_name=sheet)
cur_dict = {}
for column in df.columns:
cur_dict[column] = df[column].tolist()
print(cur_dict.keys())
dict[sheet] = cur_dict
print(dict.keys())
return dict
def get_sheet_list(dict, sheets=None, key="url"):
images_list = [dict[key]] if sheets is None else [dict[sheet_name][key] for sheet_name in sheets]
images_full = []
for images, sheet in zip(images_list, sheets):
print(f"{sheet}: {len(images)}")
images_full = images_full + images
return images_full
def test_load_save_obj():
image_name = "000000243355_zebra"
obj = f"./unit_test/{image_name}.obj"
obj = load_obj(obj)
export_obj(f"./unit_test/{image_name}_resave", obj)
if __name__ == '__main__':
# test = [(1,2), (3,4)]
# dump_pkl('test.pkl', test)
# print(load_pkl('test.pkl'))
# xyz = np.random.rand(1000, 3)
# write_pointcloud("test.ply", xyz)
# xyz = np.random.rand(1000, 3)
# write_pointcloud("test.ply", xyz)
# pass
# test_depth_8uc3_encode()
# test_zbuffer_to_depth()
# fuse_frames_of_depth_video()
test_load_save_obj()