Permalink
Browse files

Commenting is done.

  • Loading branch information...
fregu856
fregu856 committed Sep 27, 2017
1 parent 6aace6e commit f095e346037c03a50d5e0c3c4baea50b3c402332
Showing with 17 additions and 10 deletions.
  1. +17 −10 run_on_sequence.py
View
@@ -1,8 +1,5 @@
import numpy as np
import cPickle
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import tensorflow as tf
import cv2
import os
@@ -11,23 +8,26 @@
from model import ENet_model
#project_dir = "/home/fregu856/segmentation/"
project_dir = "/root/segmentation/"
data_dir = project_dir + "data/"
model_id = "test" # (change this to not overwrite all log data when you train the model)
model_id = "sequence_run"
batch_size = 4
img_height = 512
img_width = 1024
model = ENet_model(model_id, img_height=img_height, img_width=img_width, batch_size=batch_size)
model = ENet_model(model_id, img_height=img_height, img_width=img_width,
batch_size=batch_size)
no_of_classes = model.no_of_classes
# load the mean color channels of the train imgs:
train_mean_channels = cPickle.load(open("data/mean_channels.pkl"))
# load the sequence data:
seq_frames_dir = "/root/data/cityscapes/leftImg8bit/demoVideo/stuttgart_02/"
seq_frame_paths = []
frame_names = sorted(os.listdir(seq_frames_dir))
for step, frame_name in enumerate(frame_names):
@@ -41,6 +41,7 @@
no_of_frames = len(seq_frame_paths)
no_of_batches = int(no_of_frames/batch_size)
# define where to place the resulting images:
results_dir = model.project_dir + "results_on_seq/"
# create a saver for restoring variables/parameters:
@@ -63,19 +64,23 @@
img_path = seq_frame_paths[batch_pointer + i]
img_paths.append(img_path)
# read the image:
img = cv2.imread(img_path, -1)
img = cv2.resize(img, (img_width, img_height))
img = img - train_mean_channels
batch_imgs[i] = img
batch_pointer += batch_size
batch_feed_dict = model.create_feed_dict(imgs_batch=batch_imgs,
early_drop_prob=0.0, late_drop_prob=0.0)
# run a forward pass and get the logits:
logits = sess.run(model.logits, feed_dict=batch_feed_dict)
print "step: %d/%d" % (step+1, no_of_batches)
# save all predicted label images overlayed on the input frames to results_dir:
predictions = np.argmax(logits, axis=3)
for i in range(batch_size):
pred_img = predictions[i]
@@ -87,12 +92,14 @@
img_name = img_file_name.split(".png")[0]
pred_path = results_dir + img_name + "_pred.png"
merged_img = 0.3*img + 0.7*pred_img_color
overlayed_img = 0.3*img + 0.7*pred_img_color
cv2.imwrite(pred_path, merged_img)
cv2.imwrite(pred_path, overlayed_img)
# create a video of all the resulting overlayed images:
fourcc = cv2.cv.CV_FOURCC("M", "J", "P", "G")
out = cv2.VideoWriter(results_dir + "cityscapes_stuttgart_02_pred.avi", fourcc, 20.0, (img_width, img_height))
out = cv2.VideoWriter(results_dir + "cityscapes_stuttgart_02_pred.avi", fourcc,
20.0, (img_width, img_height))
frame_names = sorted(os.listdir(results_dir))
for step, frame_name in enumerate(frame_names):

0 comments on commit f095e34

Please sign in to comment.