def test_ScoreSchema_addObj_number_type():
# make sure all data gets converted to int-type
# addCircle
inp_data = [11.0, 12, float(15), float(0)]
ss = ScoreSchema()
ss.addCircle(inp_data, objEnum=0)
ret_data = ss.getAll()['0']['data']
assert all(map(lambda elem: type(elem) is int, ret_data))
assert ret_data == [11, 12, 15, 0]
# addRay
inp_data = [[11.0, 12.9999], [float(15), float(0.1)]]
ss = ScoreSchema()
ss.addRay(inp_data, objEnum=0)
ret_data = ss.getAll()['0']['data']
assert all(map(lambda elem: type(elem) is int, ret_data[0]))
assert all(map(lambda elem: type(elem) is int, ret_data[1]))
assert ret_data == [[11, 12], [15, 0]]
class TrackFactory:
''' handles tracking processing and tracking-stored-data '''
def __init__(self, on=False):
self.on = on
self.currentFrame = None
self.currentFrameInd = -1
self.trackingOnPrevious = None
self.bTrackingOnChange = False
self.currentTrackSuccess = False
self.currentTrackScore = ScoreSchema()
self.threshInitial = [(0, 0, 0), (255, 255, 255)]
self.threshes = []
self.lastFrame = None
self.declaredBallColor = ""
self.bPerformTrainOnNewData = False
self.trainingData = []
self.trainingThreshes = []
self.bTrackTimer = False
self.trackTimerData = {}
self.savedParams = None
# tp_: tracking parameters
self.tp_trackAlgoEnum = 0
self.tp_tracking_blur = 1
self.tp_repair_iterations = 1
self.tp_b_hsv = False
# TrackingAlgo class inherits TrackingTemplate
# and is instantiated here?
def setInit(self, ballColor=""):
self.declaredBallColor = ballColor
if self.declaredBallColor == "green":
self.threshInitial = [(29, 86, 6), (64, 255, 255)]
self.threshes.append(tuple(self.threshInitial))
self.threshes.append(((20, 60, 6), (40, 255, 255)))
if self.declaredBallColor == "orange":
thresh1 = ((6, 30, 120), (64, 255, 255))
thresh2 = ((64, 100, 180), (90, 255, 255))
thresh3 = ((90, 120, 200), (120, 255, 255))
self.threshInitial = thresh1
self.threshes = [thresh1, thresh2, thresh3]
def setAlgoEnum(self, algoEnum):
self.tp_trackAlgoEnum = algoEnum
def setCmd(self,
trackingOn,
outputParams=None,
alterParams=None,
resetParams=None):
self.on = trackingOn
# this is for forcing a display "redraw" when track toggles On/Off
if self.trackingOnPrevious is not None:
if trackingOn != self.trackingOnPrevious:
self.bTrackingOnChange = True
else:
self.bTrackingOnChange = False
self.trackingOnPrevious = trackingOn
if outputParams is not None:
if outputParams:
self.outputParams()
g.switchOutputParams = False
if alterParams is not None:
if alterParams:
self.saveParams()
self.alterParams()
g.switchAlterParams = False
self.bTrackingOnChange = True
if resetParams is not None:
if resetParams:
self.restParams()
g.switchResetParams = False
self.bTrackingOnChange = True
def outputParams(self):
try:
with open("notes/tracking_params.json", "w") as f:
json.dump(self.getTrackParams(), f, indent=4)
except:
print 'failed to output tracking params'
def alterParams(self):
try:
with open("notes/tracking_params.json", "r") as f:
d_params = json.load(f)
except Exception as e:
print 'failed to read in json'
print e.message
try:
self.setTrackParams(**d_params)
except Exception as e:
print 'failed to set track params from json'
print e
def saveParams(self):
self.savedParams = self.getTrackParams()
def restParams(self):
self.setTrackParams(**self.savedParams)
def resetTracker(self):
self.trackTimerData = {}
def getTrackOnChange(self):
''' return True to make a pass thru inner loop of guiview
only updating state related to tracking. This allows us to make
changes immediately apparent in Display'''
if self.bTrackingOnChange is None:
return False
else:
return self.bTrackingOnChange
def getTrackScore(self):
if not (self.on): return None
return self.currentTrackScore.getAll()
def setFrameInd(self, frameInd):
self.currentFrameInd = frameInd
def setFrame(self, currentFrame):
if not (self.on): return
self.currentFrame = currentFrame
def setFrameScore(self, frameScoreData):
if not (self.on): return
if frameScoreData is None: return
if len(frameScoreData) != 2: return
frameType, frameScore = frameScoreData
objScoring = ScoreSchema()
objScoring.load(frameScore)
circleDataObj0 = objScoring.getData(objEnum=0)
if frameType == "training":
datum = self.buildTrainingDatum(circleDataObj0, self.currentFrame)
if datum is not None:
self.trainingData.append(datum)
if self.bPerformTrainOnNewData:
self.trainProc()
if frameType == "scoring":
pass
#do evaluation
@classmethod
def buildTrainingDatum(cls, cropRect, img):
if cropRect is None or img is None:
return None
try:
datum = {}
datum['cropRect'] = cropRect
datum['cropImg'] = crop_img(img.copy(), cls.absRect(cropRect))
return datum
except:
return None
def trainProc(self):
''' take training data and build thresh hi/lo from them '''
if len(self.trainingData) < 1: return
# only run iterthresh on latest img
img = self.trainingData[len(self.trainingData) - 1].get(
'cropImg', None)
if img is None: return
if img.shape[0] < 1 or img.shape[1] < 1: return
circle_img = pixlist_to_pseduoimg(filter_pixels_circle(img))
out_thresh = iterThreshA(circle_img, goal_pct=.95, steep=False)
_lo, _hi = out_thresh[1][3], out_thresh[1][4]
self.trainingThreshes.append((_lo, _hi))
self.combine_threshes(self.trainingThreshes)
self.threshInitial = (tuple(map(int, _lo)), tuple(map(int, _hi)))
#TODO - update self.threshes, if necessary
def setTrackTimer(self, bTrackTimer):
if isinstance(bTrackTimer, bool):
self.bTrackTimer = bTrackTimer
def getTrackTimerData(self):
return self.trackTimerData
def getTrackTimerDataCurrent(self):
return self.trackTimerData.get(self.currentFrameInd - 1, -1)
def setTrackParams(self,
tracking_blur=None,
repair_iterations=None,
thresh_lo=None,
thresh_hi=None,
threshes=None):
if tracking_blur is not None:
self.tp_tracking_blur = tracking_blur
if repair_iterations is not None:
self.tp_repair_iterations = repair_iterations
if thresh_lo is not None:
self.threshInitial[0] = tuple(thresh_lo)
if thresh_hi is not None:
self.threshInitial[1] = tuple(thresh_hi)
if threshes is not None:
self.threshes = copy.copy(threshes)
def getTrackParams(self):
params = {}
params['tracking_blur'] = self.tp_tracking_blur
params['repair_iterations'] = self.tp_repair_iterations
params['thresh_lo'] = self.threshInitial[0]
params['thresh_hi'] = self.threshInitial[1]
params['threshes'] = copy.copy(self.threshes)
return params
def trackFrame(self, b_log=False):
'''
Wrapper function for a particular trackAlgo:
- unpack parameters
- select the particular track algo to run
- [possibly] time the function (if self.bTrackTimer)
- [possibly] return log of img transform steps (if b_log)
questions/todos:
[ ] pass in objEnum to trackAlgo for multi-obj tracking
'''
if not (self.on): return
tracking_blur = self.tp_tracking_blur
repair_iterations = self.tp_repair_iterations
thresh_lo = self.threshInitial[0]
thresh_hi = self.threshInitial[1]
threshes = copy.copy(self.threshes)
last_frame = self.lastFrame
if self.bTrackTimer:
t0 = time.time()
ret = None
if self.tp_trackAlgoEnum == 0:
ret = self.trackDefault(
tracking_blur=tracking_blur,
repair_iterations=repair_iterations,
thresh_lo=thresh_lo,
thresh_hi=thresh_hi,
b_log=b_log,
objEnum=0 #TODO-SS
)
elif self.tp_trackAlgoEnum == 1:
ret = self.trackDemoNew(tracking_blur=tracking_blur,
repair_iterations=repair_iterations,
thresh_lo=thresh_lo,
thresh_hi=thresh_hi,
b_log=b_log
# ,objEnum = 0
)
elif self.tp_trackAlgoEnum == 2:
ret = self.trackMultiThresh1(tracking_blur=tracking_blur,
repair_iterations=repair_iterations,
threshes=threshes,
b_log=b_log
# ,objEnum = 0
)
elif self.tp_trackAlgoEnum == 3:
ret = self.trackDummyRichLog(tracking_blur=tracking_blur,
repair_iterations=repair_iterations,
threshes=threshes,
b_log=b_log,
last_frame=last_frame
# ,objEnum = 0
)
else:
print 'trackAlgoEnum not recognized'
if self.bTrackTimer:
if self.currentFrameInd not in self.trackTimerData.keys():
t_proc = time.time() - t0
self.trackTimerData[self.currentFrameInd] = t_proc
if ret is not None:
return ret
# tracker algos --------
def trackDefault(self,
tracking_blur,
repair_iterations,
thresh_lo,
thresh_hi,
b_log=False,
objEnum=0):
'''
trackDefault:
Template for writing a track algo.
- add documentation notes, describing how this algo is different;
in this case we're describing the templating.
- organize all parameters used in func args; these are
retreived from the instance in parent function, trackFrame,
and thus read-only.
- add b_log and objEnum defaults to args
- write trackAlgo output data to instance properties:
self.currentTrackSuccess
self.currentTrackScore (as a DataSchema.ScoreSchema)
- return None, unless b_log
- in which case, include a b_log section:
- add all possible transforms to 'keys' list.
- set 'data' for each possible key, mimicing control flow
for early return in the function
(this will be used to debug in notebooks, but is
not necessary for most purposes.)
questions / todos:
[ ] will we overwrite data before it hits b_log return data?
'''
img_t = transformA(self.currentFrame.copy(), tracking_blur)
img_mask = threshA(img_t, threshLo=thresh_lo, threshHi=thresh_hi)
if not (img_mask is None) and (img_mask.sum() != 0):
img_mask_2 = repairA(img_mask, iterations=repair_iterations)
x, y = find_xy(img_mask_2)
radius = find_radius(img_mask_2)
if radius > 0:
self.currentTrackSuccess = True
self.currentTrackScore.addCircle(self.circleToRect(
(x, y, radius)),
objEnum=objEnum)
else:
self.currentTrackSuccess = False
self.currentTrackScore.reset()
else:
self.currentTrackSuccess = False
self.currentTrackScore.reset()
if b_log:
keys = [
'img_t', 'img_mask', 'img_repair', 'xy', 'radius',
'scoreCircle'
]
data = OrderedDict()
for k in keys:
data[k] = None
data['img_t'] = img_t
data['img_mask'] = img_mask
if img_mask is not None:
data['img_repair'] = img_mask_2
data['xy'] = (x, y)
data['radius'] = radius
data['scoreCircle'] = self.currentTrackScore.getObjRect(0)
return data
def trackDemoNew(self,
tracking_blur,
repair_iterations,
thresh_lo,
thresh_hi,
b_log=False):
'''
trackDemoNew:
Example for adding a non-default trackAlgo.
Only a few small changes to this function:
- allow us to bypass repairA step by setting
repair_iterations to 0
- early return when img_mask is blank
questions / todos:
[x] verify early return skips a section and verify the notebook
workflow process handles the missing data gracefully
in mutliPlot()
[ ] is a tracking_blur = 1 do any changes?
[ ] does repair_iteration = 0 work in trackDefault()?
'''
img_t = transformA(self.currentFrame.copy(), tracking_blur)
img_mask = threshA(img_t, threshLo=thresh_lo, threshHi=thresh_hi)
if img_mask.sum() != 0:
if repair_iterations > 0:
img_repair = repairA(img_mask, iterations=repair_iterations)
img_terminal = img_repair
else:
img_terminal = img_mask
x, y = find_xy(img_terminal)
radius = find_radius(img_terminal)
if radius > 0:
self.currentTrackSuccess = True
self.currentTrackScore.addCircle(self.circleToRect(
(x, y, radius)),
objEnum=0)
else:
self.currentTrackSuccess = False
self.currentTrackScore.reset()
else:
self.currentTrackSuccess = False
self.currentTrackScore.reset()
if b_log:
keys = [
'img_t', 'img_mask', 'img_repair', 'img_terminal', 'xy',
'radius', 'scoreCircle'
]
data = OrderedDict()
for k in keys:
data[k] = None
data['img_t'] = img_t
data['img_mask'] = img_mask
if img_mask.sum() != 0:
data['img_terminal'] = img_terminal
data['xy'] = (x, y)
data['radius'] = radius
if repair_iterations > 0:
data['img_repair'] = img_repair
if radius > 0:
data['scoreCircle'] = self.currentTrackScore.getObjRect(0)
return data
def trackMultiThresh1(self,
tracking_blur,
repair_iterations,
threshes,
b_log=False):
'''
trackMultiThresh1:
Evolved from trackDemoNew; uses multiple thresh intervals
questions / todos:
[ ] multiple threshes
'''
img_t = transformA(self.currentFrame.copy(), tracking_blur)
img_mask = threshMultiOr(img_t, threshes=threshes)
if img_mask.sum() != 0:
if repair_iterations > 0:
img_repair = repairA(img_mask, iterations=repair_iterations)
img_terminal = img_repair
else:
img_terminal = img_mask
x, y = find_xy(img_terminal)
radius = find_radius(img_terminal)
if radius > 0:
self.currentTrackSuccess = True
self.currentTrackScore.addCircle(self.circleToRect(
(x, y, radius)),
objEnum=0)
else:
self.currentTrackSuccess = False
self.currentTrackScore.reset()
else:
self.currentTrackSuccess = False
self.currentTrackScore.reset()
if b_log:
keys = [
'img_t', 'img_mask', 'img_repair', 'img_terminal', 'xy',
'radius', 'scoreCircle'
]
data = OrderedDict()
for k in keys:
data[k] = None
data['img_t'] = img_t
data['img_mask'] = img_mask
if img_mask.sum() != 0:
data['img_terminal'] = img_terminal
data['xy'] = (x, y)
data['radius'] = radius
if repair_iterations > 0:
data['img_repair'] = img_repair
if radius > 0:
data['scoreCircle'] = self.currentTrackScore.getObjRect(0)
return data
def trackDummyRichLog(self,
tracking_blur,
repair_iterations,
threshes,
last_frame,
b_log=False):
'''
trackDummyRichLog:
copied from trackMultiThresh1; here we're just adding
extra track_log keys like img_diff and img_diff_repair
to see how they'll impact reporting functions downstream.
this is the start though of looking at adding a movement
element to trackers and we're adding the last_frame input /
self.lastFrame output.
'''
img_t = transformA(self.currentFrame.copy(), tracking_blur)
img_mask = threshMultiOr(img_t, threshes=threshes)
# add new-lines in here:
if last_frame is not None and False:
img_diff = cv2.absdiff(self.currentFrame, last_frame)
if repair_iterations > 0:
img_diff_repair = repairA(img_diff,
iterations=repair_iterations)
self.lastFrame = self.currentFrame.copy()
if img_mask.sum() != 0:
if repair_iterations > 0:
img_repair = repairA(img_mask, iterations=repair_iterations)
img_terminal = img_repair
else:
img_terminal = img_mask
x, y = find_xy(img_terminal)
radius = find_radius(img_terminal)
if radius > 0:
self.currentTrackSuccess = True
self.currentTrackScore.addCircle(self.circleToRect(
(x, y, radius)),
objEnum=0)
else:
self.currentTrackSuccess = False
self.currentTrackScore.reset()
else:
self.currentTrackSuccess = False
self.currentTrackScore.reset()
if b_log:
keys = [
'img_t', 'img_mask', 'img_repair', 'img_dummy', 'img_dummy_2',
'img_diff', 'img_diff_repair'
'img_terminal', 'img_terminal_2'
'xy', 'radius', 'scoreCircle'
]
data = OrderedDict()
for k in keys:
data[k] = None
data['img_t'] = img_t
data['img_mask'] = img_mask
data['img_dummy'] = transformA(self.currentFrame.copy(), 11)
data['img_dummy_2'] = transformA(self.currentFrame.copy(), 5)
if last_frame is not None and False:
data['img_diff'] = img_diff
data['img_diff_repair'] = img_diff_repair
if img_mask.sum() != 0:
data['img_terminal'] = img_terminal
data['img_terminal_2'] = img_terminal
data['xy'] = (x, y)
data['radius'] = radius
if repair_iterations > 0:
data['img_repair'] = img_repair
if radius > 0:
data['scoreCircle'] = self.currentTrackScore.getObjRect(0)
return data
# helper functions ------
@staticmethod
def absRect(input_rect):
''' takes an (opencv style) relative rect, returns an absolute rect.
(x0,y0, d_x, d_y) -> ((xo,y0),(x1, y1))
note: must be tuples, not lists; to use in opencv functions
'''
x = copy.copy(input_rect)
rect = ((int(x[0]), int(x[1])), (int(x[0] + x[2]), int(x[1] + x[3])))
return rect
@staticmethod
def circleToRect(input_circle):
''' takes x,y, radius, fits to enclosing relative-format rect
(x,y, radius) -> (x0,y0, d_x, d_y)
'''
x, y, radius = copy.copy(input_circle)
x0 = int(x - radius)
y0 = int(y - radius)
dx = int(2 * radius)
dy = int(2 * radius)
return (x0, y0, dx, dy)
@staticmethod
def combine_threshes(data, liberal=True):
''' data is a list of (lo, hi) 3-ple's; find the union '''
if len(data) < 1:
return (np.array([0, 0, 0], dtype='uint8'),
np.array([255, 255, 255], dtype='uint8'))
_lo, _hi = [[255, 255, 255], [0, 0, 0]]
for row in data:
lo, hi = row[0], row[1]
for i, clr in enumerate(lo):
if clr < _lo[i]:
_lo[i] = clr
for i, clr in enumerate(hi):
if clr > _hi[i]:
_hi[i] = clr
return (np.array(_lo, dtype='uint8'), np.array(_hi, dtype='uint8'))
def show_scoring_on_off_1(input_test_child_dir,
input_circle_data,
b_rebench=False):
''' test that turning show_scoring on/off affects main display panel
test that show_scoring=on + no scoring data is handled
input params: - tests have dif size frames
- tests have dif scoring-data
'''
# setup ------
TEST_CHILD_DIR = input_test_child_dir
stub_frame = cv2.imread(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR, "stubframe.png"))
bench_yes_scoring = cv2.imread(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR, "bench_yes_score.png"))
bench_no_scoring = cv2.imread(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR, "bench_no_score.png"))
some_scoring = ScoreSchema()
some_scoring.addCircle(input_circle_data)
stub_some_score = some_scoring.getAll()
none_scoring = ScoreSchema()
stub_none_score = none_scoring.getAll()
diff = ImgDiff(log_path=DIFF_LOG_DIR)
# run test -----
stage = StagingDisplay()
stage.all_display_methods(b_showscoring=True,
stub_frame=stub_frame.copy(),
stub_scorecurrent=copy.deepcopy(stub_some_score))
scoring_on_output = stage.mock_get_frame()
stage = StagingDisplay()
stage.all_display_methods(
b_showscoring=False #test-variable
,
stub_frame=stub_frame.copy(),
stub_scorecurrent=copy.deepcopy(stub_some_score))
scoring_off_output = stage.mock_get_frame()
stage = StagingDisplay()
stage.all_display_methods(
b_showscoring=True,
stub_frame=stub_frame.copy(),
stub_scorecurrent=copy.deepcopy(stub_none_score) #test-variable
)
scoring_none_output = stage.mock_get_frame()
#rebench ---
if b_rebench:
if verifyAction(prefix="\nrebench:" + input_test_child_dir):
return
cv2.imwrite(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR,
"bench_yes_score.png"), scoring_on_output)
cv2.imwrite(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR,
"bench_no_score.png"), scoring_off_output)
return
# verify ----
assert diff.diffImgs(bench_yes_scoring, scoring_on_output)
assert diff.diffImgs(bench_no_scoring, scoring_off_output)
assert diff.diffImgs(bench_no_scoring, scoring_none_output)
assert diff.diffImgs(scoring_on_output, scoring_off_output,
noLog=True) == False
def show_tracking_on_off_1(input_test_child_dir,
input_circle_data,
b_rebench=False):
''' test that turning tracking on/off affecta main_display
and that it affects score_display
input params: - tests have dif size frames
- tests have dif scoring-data
TODO:
[x] modulo zero resize
'''
# setup ------
TEST_CHILD_DIR = input_test_child_dir
stub_frame = cv2.imread(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR, "stubframe.png"))
bench_yes_tracking = cv2.imread(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR,
"bench_yes_track_main.png"))
bench_no_tracking = cv2.imread(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR,
"bench_no_track_main.png"))
bench_score = cv2.imread(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR,
"bench_yes_track_score.png"))
some_scoring = ScoreSchema()
some_scoring.addCircle(input_circle_data)
stub_some_score = some_scoring.getAll()
none_scoring = ScoreSchema()
stub_none_score = none_scoring.getAll()
diff = ImgDiff(log_path=DIFF_LOG_DIR)
# run test -----
#1:
stage = StagingDisplay()
stage.all_display_methods(stub_frame=stub_frame.copy(),
stub_trackscore=copy.deepcopy(stub_some_score))
main1 = stage.mock_get_frame()
score1 = stage.mock_get_score_frame()
#2:
stage = StagingDisplay()
stage.all_display_methods(
stub_frame=stub_frame.copy(),
stub_trackscore=copy.deepcopy(stub_none_score) #test-variable
)
main2 = stage.mock_get_frame()
score2 = stage.mock_get_score_frame()
#3:
stage = StagingDisplay()
stage.all_display_methods(
b_showscoring=True #test-variable
,
stub_frame=stub_frame.copy(),
stub_trackscore=copy.deepcopy(stub_some_score))
main3 = stage.mock_get_frame()
score3 = stage.mock_get_score_frame()
#rebench ---
if b_rebench:
if verifyAction(prefix="\nrebench:" + input_test_child_dir):
return
cv2.imwrite(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR,
"bench_yes_track_main.png"), main1)
cv2.imwrite(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR,
"bench_no_track_main.png"), main2)
cv2.imwrite(
os.path.join(TEST_PARENT_DIR, TEST_CHILD_DIR,
"bench_yes_track_score.png"), score3)
return
# verify ----
assert diff.diffImgs(bench_yes_tracking, main1)
assert diff.diffImgs(bench_no_tracking, main2)
assert diff.diffImgs(bench_yes_tracking, main3)
assert diff.diffImgs(main1, main2, noLog=True) == False
assert diff.diffImgs(bench_score, score1)
assert score2 is None
assert diff.diffImgs(bench_score, score3)