run_lpo.py

"""
Make sure you have LPO_DIR/src in your python path
"""
 
import lpo
from lpo import *
from util import *
import cv2
import numpy as np
import os
from pylab import *
 
# ------------- filenames -----------------
dataset = "50Salads"
filename = "rgb-01-1.avi"
# ------------- Folder names -----------------
# url_files = os.path.expanduser("~/Desktop/tmp2/".format(dataset))
url_files = os.path.expanduser("~/Data/{}/raw/rgb/".format(dataset))
url_img_save = os.path.expanduser("~/Data/{}/viz/lpo/".format(dataset))
url_bbox_save = os.path.expanduser("~/Data/{}/features/lpo/S1_Cheese_C1/".format(dataset))
# I can only get the library to read images from disk (instead of from video stream)
# So save video file to temporary image and then load...
tmp_filename = os.path.expanduser("~/Desktop/tmp.png")
# ------------- User params -----------------
save_file = [False, True][0]
n_superpixels = 200
rez = (854, 480)
# I threshold out boxes that are too big or too small
min_box_len = 50
max_box_len = 250
# This list is all available overlap scores.
overlap_score = [.01, .02, .03, .05, .1, .2][-1]
# ---------------------------------------------
 
# Load the ensemble
prop = proposals.LPO()
dir_lpo = os.path.dirname(lpo.__file__).strip("src")
prop.load(dir_lpo+"/models/lpo_VOC_{}.dat".format(overlap_score))
detector = getDetector('mssf')
 
# Check that the save folder exists
if not os.path.exists(url_img_save):
    os.mkdir(url_img_save)
 
vid = cv2.VideoCapture(os.path.expanduser(url_files+filename))
 
# Evaluate on the test set
i=0
while 1:
    ret, im_cv = vid.read()
    if not ret:
        break
    # im_cv = cv2.resize(im, rez)
 
    cv2.imwrite(tmp_filename, im_cv)
    im = imgproc.imread(tmp_filename)
 
    # Extract object proposals from oversegmented regions
    over_segs = segmentation.generateGeodesicKMeans( detector, [im], n_superpixels)
    props = prop.propose(over_segs)[0]
    segs = props[0].s
 
    # Display bounding boxes
    boxes = props[0].toBoxes()
    for b in boxes:
        x1,y1,x2,y2 = b
        if (min_box_len < (x2 - x1) < max_box_len) and \
            (min_box_len < (y2 - y1) < max_box_len):
            cv2.rectangle(im_cv,(x1,y1),(x2,y2),(0,255,0),1)
    cv2.imshow("im", im_cv)
    cv2.waitKey(30)
 
    # # Overlay superpixels onto image to create heatmap
    # im_ = np.zeros_like(im_cv[:,:,2])*1.
    # for o in props[0].p:
    #     for spx in np.nonzero(o)[0]:
    #         im_[segs==spx] += 1
    #
    #     # im_ += (segs == np.nonzero(o)[0][:,None,None]).sum(0)
    #
    # cv2.imshow("heatmap", im_/im_.max())
    # cv2.waitKey(30)
 
    # Save image and bounding box/segmentation to file
    if save_file:
        base_filename = str(i*framerate)+".jpg"
        cv2.imwrite(url_img_save+base_filename, im_cv)
        saveProposalsHDF5( props, url_bbox_save+base_filename+'.hf5', True, True )
 
    i+=1