def select(exists_path, pool_path, div=200):
# 先计算output_path的已有图像
import cv
import dfl
import random
width = 800
trans = cv.trans_fn(-1, 1, 0, width)
img = cv.cv_new((width, width))
for f in io.progress_bar_generator(os.listdir(exists_path),
"Existing Imgs"):
if f.endswith(".png") or f.endswith("jpg"):
img_path = os.path.join(exists_path, f)
dfl_img = dfl.dfl_load_img(img_path)
pitch, yaw, _ = dfl.dfl_estimate_pitch_yaw_roll(dfl_img)
pitch = trans(pitch)
yaw = trans(yaw)
cv.cv_circle(img, (pitch, yaw), (128, 128, 128), width / div, -1)
time_str = get_time_str()
import shutil
pool_files = list(os.listdir(pool_path))
# random.shuffle(pool_files)
count = 0
for f in io.progress_bar_generator(pool_files,
os.path.basename(pool_path)):
if f.endswith(".png") or f.endswith(".jpg"):
img_path = os.path.join(pool_path, f)
dfl_img = dfl.dfl_load_img(img_path)
pitch, yaw, _ = dfl.dfl_estimate_pitch_yaw_roll(dfl_img)
pitch = trans(pitch)
yaw = trans(yaw)
if sum(img[yaw][pitch]) == 255 * 3:
dst = os.path.join(exists_path, "%s_%s" % (time_str, f))
shutil.copy(img_path, dst)
count += 1
cv.cv_circle(img, (pitch, yaw), (0xcc, 0x66, 0x33),
width / div, -1)
cv.cv_save(img, os.path.join(exists_path, "_select.bmp"))
io.log_info("Copy %d, Total %d" % (count, len(pool_files)))
def reload_src():
nonlocal src_img_list
nonlocal src_cur_list
src_img_list = []
if src_path:
for f in io.progress_bar_generator(os.listdir(src_path),
"Loading"):
if f.endswith(".jpg") or f.endswith(".png"):
fpath = os.path.join(src_path, f)
dfl_img = dfl.dfl_load_img(fpath)
p, y, _ = dfl.dfl_estimate_pitch_yaw_roll(dfl_img)
src_img_list.append([fno, p, y])
src_img_list.append([fno, p, -y])
src_img_list = np.array(src_img_list, "float")
src_cur_list = src_img_list
def manual_select(input_path, src_path=None):
import cv
import colorsys
import cv2
img_list = []
src_img_list = []
width = 800
ratio = 0.8
for f in io.progress_bar_generator(os.listdir(input_path), "Loading"):
if f.endswith(".jpg") or f.endswith(".png"):
fpath = os.path.join(input_path, f)
dfl_img = dfl.dfl_load_img(fpath)
p, y, _ = dfl.dfl_estimate_pitch_yaw_roll(dfl_img)
fno = int(f.split(".")[0])
img_list.append([fno, p, y])
# for i in range(10000):
# img_list.append([i,
# random.random() * 2 - 1,
# random.random() * 2 - 1])
src_img_list = []
src_cur_list = []
img_list = np.array(img_list, "float")
cur_list = img_list
src_r = width / 100 * 2.5
redius = width / 100 * 2
trans_pitch_to_x = cv.trans_fn(-1, 1, 0, width)
trans_yaw_to_y = cv.trans_fn(-1, 1, 0, width)
trans_x_to_pitch = cv.trans_fn(0, width, -1, 1)
trans_y_to_yaw = cv.trans_fn(0, width, -1, 1)
trans_r = cv.trans_fn(0, width, 0, 2)
cur_pitch_yaw = img_list[:, 1:3]
img = cv.cv_new((width, width))
cur_w = 2
cur_mid = (0, 0)
def reload_src():
nonlocal src_img_list
nonlocal src_cur_list
src_img_list = []
if src_path:
for f in io.progress_bar_generator(os.listdir(src_path),
"Loading"):
if f.endswith(".jpg") or f.endswith(".png"):
fpath = os.path.join(src_path, f)
dfl_img = dfl.dfl_load_img(fpath)
p, y, _ = dfl.dfl_estimate_pitch_yaw_roll(dfl_img)
src_img_list.append([fno, p, y])
src_img_list.append([fno, p, -y])
src_img_list = np.array(src_img_list, "float")
src_cur_list = src_img_list
def repaint():
nonlocal trans_pitch_to_x
nonlocal trans_yaw_to_y
nonlocal trans_x_to_pitch
nonlocal trans_y_to_yaw
nonlocal trans_r
nonlocal src_cur_list
nonlocal cur_list
nonlocal cur_pitch_yaw
nonlocal img
nonlocal src_path
mid = cur_mid
w = cur_w
sx = mid[0] - w / 2
sy = mid[1] - w / 2
ex = mid[0] + w / 2
ey = mid[1] + w / 2
idxs = (img_list[:, 1] >= sx) & \
(img_list[:, 2] >= sy) & \
(img_list[:, 1] <= ex) & \
(img_list[:, 2] <= ey)
cur_list = img_list[idxs]
cur_pitch_yaw = cur_list[:, 1:3]
if len(src_img_list) == 0:
src_cur_list = []
elif src_path:
idxs = (src_img_list[:, 1] >= sx) & \
(src_img_list[:, 2] >= sy) & \
(src_img_list[:, 1] <= ex) & \
(src_img_list[:, 2] <= ey)
src_cur_list = src_img_list[idxs]
trans_pitch_to_x = cv.trans_fn(sx, ex, 0, width)
trans_yaw_to_y = cv.trans_fn(sy, ey, 0, width)
trans_x_to_pitch = cv.trans_fn(0, width, sx, ex)
trans_y_to_yaw = cv.trans_fn(0, width, sy, ey)
trans_r = cv.trans_fn(0, width, 0, w)
img = cv.cv_new((width, width))
min_fno = int(cur_list[0][0])
max_fno = int(cur_list[-1][0])
trans_color = cv.trans_fn(min_fno, max_fno, 0, 1)
for _, p, y in src_cur_list:
cv.cv_point(img, (trans_pitch_to_x(p), trans_yaw_to_y(y)),
(192, 192, 192), src_r * 2 / cur_w)
for f, p, y in cur_list:
fno = int(f)
h = trans_color(fno)
s = 1
v = 1
r, g, b = colorsys.hsv_to_rgb(h, s, v)
cv.cv_point(img, (trans_pitch_to_x(p), trans_yaw_to_y(y)),
(b * 255, g * 255, r * 255), 2)
cv2.imshow("select", img)
def mouse_callback(event, x, y, flags, param):
nonlocal cur_mid
nonlocal cur_w
x = trans_x_to_pitch(x)
y = trans_y_to_yaw(y)
if event == cv2.EVENT_LBUTTONDOWN:
tr = trans_r(redius)
point = np.array([[x, y]] * len(cur_pitch_yaw), "float")
dist = np.linalg.norm(cur_pitch_yaw - point, axis=1)
idxs = dist <= tr
for f, _, _ in cur_list[idxs]:
print(f)
pass
print("-----------------------------------------")
elif event == cv2.EVENT_RBUTTONDOWN:
cur_mid = (x, y)
cur_w = cur_w * ratio
repaint()
elif event == cv2.EVENT_MBUTTONDOWN:
cur_w = cur_w / ratio
if cur_w >= 2:
cur_w = 2
cur_mid = (0, 0)
repaint()
reload_src()
cv2.namedWindow("select")
cv2.setMouseCallback("select", mouse_callback)
while True:
repaint()
key = cv2.waitKey()
if key == 13 or key == -1:
break
elif key == 114:
reload_src()