def on_video(self, video_path):
cap = cv2.VideoCapture(video_path)
while (cap.isOpened()):
ret, frame = cap.read()
rows, cols, _ = frame.shape
prevTime = time.time()
frame_small = cv2.resize(frame, (512, 256)) / 255.0
frame_small = np.rollaxis(frame_small, axis=2, start=0)
out_x, out_y = self.detector.test_on_image(np.array([frame_small]))
vis_image = draw_points(out_x[0],
out_y[0],
frame,
scale_x=cols / 512,
scale_y=rows / 256)
curTime = time.time()
sec = curTime - prevTime
fps = 1 / (sec)
s = "FPS : " + str(fps)
cv2.putText(vis_image, s, (0, 100), cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 255, 0))
cv2.imshow('frame', vis_image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
def on_images(self, *images):
print("model path: ", self.model_path)
print("image path: ", images)
for image_path in images:
test_image = cv2.imread(image_path)
rows, cols, _ = test_image.shape
small_image = cv2.resize(test_image, (512, 256)) / 255.0
small_image = np.rollaxis(small_image, axis=2, start=0)
out_x, out_y = self.detector.test_on_image(np.array([small_image]))
vis_image = draw_points(out_x[0],
out_y[0],
test_image,
scale_x=cols / 512,
scale_y=rows / 256)
cv2.imshow("res", vis_image)
cv2.imshow("img", test_image)
cv2.waitKey(0)
def on_folder(self, path, save=False, recursive=False):
import os
if recursive:
filenames = [
os.path.join(dp, f) for dp, dn, filenames in os.walk(path)
for f in filenames if os.path.splitext(f)[1] == '.jpg'
or os.path.splitext(f)[1] == ".png"
]
else:
filenames = [
os.path.join(path, f) for f in os.listdir(path)
if f.endswith((".png", ".jpg"))
]
filenames.sort()
out_stream = None
for each in filenames:
test_image = cv2.imread(each)
rows, cols, _ = test_image.shape
if save and (out_stream is None):
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
out_stream = cv2.VideoWriter("./output.avi", fourcc, 10.0,
(cols, rows))
small_image = cv2.resize(test_image, (512, 256)) / 255.0
small_image = np.rollaxis(small_image, axis=2, start=0)
out_x, out_y = self.detector.test_on_image(np.array([small_image]))
vis_image = draw_points(out_x[0],
out_y[0],
test_image,
scale_x=cols / 512,
scale_y=rows / 256) # nbatch = 1
if save:
out_stream.write(vis_image)
cv2.imshow("img", vis_image)
if cv2.waitKey(1) == 27:
break
if save:
out_stream.release()
def on_tusimple(self):
from dataset.tusimple import DatasetTusimple
from configs.PINet import ParamsTuSimple
dataset_param = ParamsTuSimple()
resize = Resize(rows=256, cols=512)
hwc_to_chw = TransposeNumpyArray((2, 0, 1))
norm_to_1 = NormalizeInstensity()
dataset = DatasetTusimple(
root_path=dataset_param.train_root_url,
json_files=dataset_param.train_json_file,
transform=transforms.Compose([resize, norm_to_1, hwc_to_chw]),
)
for i in range(len(dataset)):
sample = dataset[i]
img_src = cv2.imread(sample["image_path"])
out_x, out_y = self.detector.test_on_image(
np.array([sample["image"]]))
vis_image = draw_points(out_x[0], out_y[0], img_src)
cv2.imshow("sample", vis_image)
cv2.imshow("original image", img_src)
if cv2.waitKey(0) == 27:
break
def on_culane(self, save=False):
from dataset.culane import DatasetCULane
from configs.PINet import ParamsCuLane
dataset_param = ParamsCuLane()
resize = Resize(rows=256, cols=512)
hwc_to_chw = TransposeNumpyArray((2, 0, 1))
norm_to_1 = NormalizeInstensity()
dataset = DatasetCULane(
root_path=dataset_param.train_root_url,
index_file=dataset_param.train_json_file,
transform=transforms.Compose([resize, norm_to_1, hwc_to_chw]),
)
if save:
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
out = cv2.VideoWriter("./output.avi", fourcc, 20.0, (1640, 590))
else:
out = None
for i in range(len(dataset)):
sample = dataset[i]
img_src = cv2.imread(sample["image_path"])
out_x, out_y = self.detector.test_on_image(
np.array([sample["image"]]))
vis_image = draw_points(out_x[0],
out_y[0],
img_src,
scale_x=1640 / 512,
scale_y=590 / 256)
if save:
out.write(vis_image)
cv2.imshow("sample", vis_image)
cv2.imshow("original image", img_src)
if cv2.waitKey(1) == 27:
break
if save:
out.release()
def test_on_image(
self,
test_images: np.ndarray,
threshold_confidence: float = 0.81
) -> Tuple[List[List[int]], List[List[int]], List[np.ndarray]]:
""" predict, then post-process
:param test_images: input image or image batch
:param threshold_confidence, if confidence of detected key points greater than threshold, then will be accepted
"""
rank = len(test_images.shape)
if rank == 3:
batch_image = np.expand_dims(test_images, 0)
elif rank == 4:
batch_image = test_images
else:
raise IndexError
# start = time.time()
result = self.predict(batch_image) # accept rank = 4 only
# end = time.time()
# print(f"predict time: {end - start} [sec]") # [second]
confidences, offsets, instances = result[
-1] # trainer use different output, compared with inference
num_batch = batch_image.shape[0]
out_x = []
out_y = []
out_images = []
for i in range(num_batch):
# test on test data set
image = deepcopy(batch_image[i])
image = np.rollaxis(image, axis=2, start=0)
image = np.rollaxis(image, axis=2, start=0) * 255.0
image = image.astype(np.uint8).copy()
confidence = confidences[i].view(
self.network_params.grid_y,
self.network_params.grid_x).cpu().data.numpy()
offset = offsets[i].cpu().data.numpy()
offset = np.rollaxis(offset, axis=2, start=0)
offset = np.rollaxis(offset, axis=2, start=0)
instance = instances[i].cpu().data.numpy()
instance = np.rollaxis(instance, axis=2, start=0)
instance = np.rollaxis(instance, axis=2, start=0)
# generate point and cluster
raw_x, raw_y = generate_result(confidence, offset, instance,
threshold_confidence)
# eliminate fewer points
in_x, in_y = eliminate_fewer_points(raw_x, raw_y)
# sort points along y
in_x, in_y = sort_along_y(in_x, in_y)
in_x, in_y = eliminate_out(in_x, in_y)
in_x, in_y = sort_along_y(in_x, in_y)
in_x, in_y = eliminate_fewer_points(in_x, in_y)
result_image = draw_points(in_x, in_y, deepcopy(image))
out_x.append(in_x)
out_y.append(in_y)
out_images.append(result_image)
return out_x, out_y, out_images