(* ::Package:: *)

$dir = FileNameJoin[{NotebookDirectory[], "Resources"}];


net := net = Import[FileNameJoin[{$dir, "YOLOv2.wlnet"}]];


labelsCOCO = Import[FileNameJoin[{$dir, "COCOclasses.json"}]];


softmax[x_]:=#/Total[Exp[x]]&/@Exp[x]
splitList[list_,lengths_]:=Take[list,{1,0}+#]&/@Partition[Prepend[Accumulate[lengths],0],2,1]

regionLayer[tensor_] := Module[
	{xxx, coords, obj, classes},
	
	xxx = Transpose[Partition[tensor, ($coord + 1 + $classes)]];
	{coords, obj, classes} = splitList[xxx, {$coord, 1, $classes}];
	
	coords[[{1,2}]] = LogisticSigmoid[coords[[{1,2}]]];
	obj = LogisticSigmoid[obj];
	classes = softmax[classes];
	(*Transpose[Join[coords, obj, classes]]*)
	{coords, obj, classes}
]

scaleImage[img_?ImageQ] := ImageCrop[ImageResize[img, {416}, Resampling -> "Linear"], {416, 416}, Padding -> .5]


anchors = {
	{0.57273, 0.677385},
	{1.87446,2.06253},
	{3.33843,5.47434},
	{7.88282,3.52778},
	{9.77052,9.16828}
};


$w = $h = 13;
$n = Length[anchors];
$classes = 80;
$coord = 4;


getRegionBox[coords_] := 
With[{grid = Table[{i, j}, {j, 0, $w - 1}, {i, 0, $h - 1}]},
	Join[
		Transpose[Transpose[coords[[{1, 2}]], {3, 4, 1, 2}] + grid, {3, 4, 1, 2}] / {$w, $h},
		Exp[coords[[{3, 4}]]]Transpose[anchors] / {$w, $h}
	]
]


(*IoU = RegionMeasure@RegionIntersection[#1,#2]/RegionMeasure@RegionUnion[#1,#2]&;*)


ClearAll[intersection, area, IoU]
IoU := With[{int = intersection[##]}, int / (area[#1] + area[#2] - int)]&
area := Times @@ Subtract @@@ Transpose[List @@ #] &
intersection := intersection =  
With[{c = Compile[
	{{box1, _Real, 2}, {box2, _Real, 2}},
	Module[{x1, y1, x2, y2, w, h},
	
		x1 = Max[box1[[1,1]], box2[[1,1]]];
		y1 = Max[box1[[1,2]], box2[[1,2]]];
		x2 = Min[box1[[2,1]], box2[[2,1]]];
		y2 = Min[box1[[2,2]], box2[[2,2]]];	
	
		w = Max[0., x2 - x1];
		h = Max[0., y2 - y1];
	
		w * h
	],
	RuntimeAttributes -> {Listable},
	Parallelization -> True,
	RuntimeOptions -> "Speed",
	CompilationTarget -> "C"
]},
	c[List @@ #1, List @@ #2]&
]


ClearAll[nonMaxSuppression]
nonMaxSuppression[overlapThreshold_][detection_] := 
Module[{boxes, confidence},
	Fold[
		{list, new} \[Function] If[NoneTrue[list, IoU[#[[2]], new[[2]]] > overlapThreshold &], Append[list, new], list],
		Sequence @@ TakeDrop[detection, 1]
	]
]


detection[img_, detectionThreshold_ : .2, overlapThreshold_ : .4] := 
Module[
	{h, w, scale, netres, coords, obj, classes, boxes, padding, classProb, probable, finals},
	
	netres = Internal`CheckImageCache[{"YOLO", Hash[scaleImage[img]]}];
	If[FailureQ[netres],
		
		{coords, obj, classes} = netres = regionLayer[net[scaleImage[img]]];
		Internal`SetImageCache[{"YOLO", Hash[scaleImage[img]]}, netres],
		
		{coords, obj, classes} = netres
	];

	{h, w} = ImageDimensions[img];
	scale = Max[{h, w}] / 416;
	padding = 416(1 - {h, w} / Max[{h, w}]) / 2;

			
	boxes = Apply[
		Rectangle[
			scale {#1 - #3/2 - padding[[1]], 416 - #2 - #4/2 - padding[[2]]},
			scale {#1 + #3/2 - padding[[1]], 416 - #2 + #4/2 - padding[[2]]}
		]&,
		416 Flatten[getRegionBox[coords], {{2,3,4}, {1}}],
	1];
	
	(*each class probability is rescaled with the box objectivness *)
	(* very small probability are thresholded *)
	classProb = Threshold[Transpose[Transpose[Flatten[classes, {{1},{2,3,4}}]] Flatten[obj]], 10^-2];
	
	probable = Position[classProb, p_ /; p > detectionThreshold];
	
	If[Length[probable] == 0, Return[{}]];
	
	finals = Join @@ Values @ GroupBy[
		MapThread[
			{labelsCOCO[[#1]], boxes[[#2]], Extract[classProb, {#1, #2}]}&,
			Transpose[probable]
		],
		First, nonMaxSuppression[overlapThreshold] @* SortBy[-Last[#]&]
	]
]


highlightDetection[img_, {}] := HighlightImage[img, {}]
highlightDetection[img_, detection_] := 
BlockRandom[ SeedRandom[0];
HighlightImage[img,

	Transpose@MapThread[
		{
			{RandomColor[], Opacity[.2],#2},
			Style[
				Inset[#1<>" ("<>ToString@Round[#3,.01]<>")",#2[[1]],{Left,Top}],
				FontSize -> Scaled[.02], GrayLevel[0,1], Background->GrayLevel[1,.7]
			]
		}&,
		Transpose[detection]
	],
	ImagePadding -> Scaled[.02]
]]
