%% calcColoniesOnPlate.m
% segments colonies on an agar plate using a watershed algorithm
% Amir Mitchell 01/01/2019

%% User definitions
strFileName = '100mmPlate_121col.jpg';
colonyAreaCutoff = [25 300]; % size range for colony area (in pixels)
circularityCutoff = [0.8 1/0.8]; % min/max ratio of colony length to equivdiamiter to qualify as a circle object
smoothDiskSize = 50;

%% open image, convert to gray, clean up, convert to BW
im0 = rgb2gray(imresize(imread(strFileName),0.5)); % open file

%% crop the image by user defined bounding box
figure; imshow(im0); hold on;

% pick xy of bounding box (lower-left and upper-right)
uiwait(msgbox({'Mark bounding box for plate (mouse left-click)';'Mark upper-left and lower-right corners'},'User input','modal'));
xy = ginput(2);
x = xy(1,1); y = xy(1,2);
w = xy(2,1)-xy(1,1); h = xy(2,2)-xy(1,2);
r = floor(min([w h])/2);
l = r*2; % length of each size of bounding box (made to by odd)
rectangle('Position',[x y l l],'Curvature',0,'EdgeColor','r');
rectangle('Position',[x y l l],'Curvature',1,'EdgeColor','r');
pause(1);

se = strel('disk',50,0); % structural element tophat
im1 = imcrop(im0,[x y l l]); % crop the image and make sure its square
im2 = imtophat(im1,se);  % correct uneven illumination  
im3 = imadjust(im2,stretchlim(im2),[]); % strech limits for min and max
im4 = medfilt2(im3); % remove salt-pepper noise 

bw1 = im2bw(im4,graythresh(im4));
se = strel('disk',r,0); % structural element an disk the will cover plate, 0 for now approximization in creating disk
bw2 = (bw1 & se.Neighborhood); % removing all pixels outside the plate

figure; 
subplot(3,1,1); imshow(im1); title('Original');
subplot(3,1,2); imshow(im2); title('illumination correction (tophat)');
subplot(3,1,3); imshow(bw2); title('Final BW');


%% Use watershed to find and separate colonies, remove tiny colonies
D = bwdist(~bw2); % make a distance matrix from BW image
%figure; imshow(D,[],'InitialMagnification','fit')
D = -D; 
D(~bw2) = Inf; % force background to be -inf distance
L = watershed(D,8); % make a label image by watershed
L(~bw2) = 0; % force label background of to be 0;

% extra step to remove tiny colonies
TF = false(size(L)); TF(L~=0)=true; % BW version of watershewd labels
TF2 = bwareaopen(TF,colonyAreaCutoff(1)); % remove tiny colonies
L(~TF2) = 0; % remove small colonies from watershewd labels
rgb = label2rgb(L,'lines','k'); % this rgb image will be used later
cc = regionprops(TF2,'All');

% count only within circle
colPositions = cat(1,cc.Centroid); % position of all legal colonies
colArea =  cat(1,cc.Area);
colLength = cat(1,cc.MajorAxisLength);
colEquivDiameter = cat(1,cc.EquivDiameter);
colCircularity = colLength./colEquivDiameter;
tfCircle = (colCircularity>circularityCutoff(1) & colCircularity<circularityCutoff(2));
tfArea = (colArea>colonyAreaCutoff(1) & colCircularity<colonyAreaCutoff(2));
truelCol = [tfCircle & tfArea]; % legal colonies by area and crcularity
falseCol = [~tfCircle | ~tfArea];
pixelToDelete =  cat(1,cc(falseCol).PixelList); % pixels of false colonies
for i=1:size(pixelToDelete,1)
    rgb(pixelToDelete(i,2),pixelToDelete(i,1),:) = 0;
end

%% make a QA figure (only detected colonies are colored)
colNum = sum(truelCol);
fh = figure;
pause(1);
h1 = imshow(im4);
hold on;
h2 = imshow(rgb);
set(h2,'AlphaData',0.5);
cc = regionprops(TF2,'All');
allArea = [cc.Area];
strTitle = [strrep(strFileName,'_','-') ' (' num2str(colNum) ' colonies)'];
title(strTitle); % title by file name, replace '_' with '-'

colX = colPositions(:,1); colY = colPositions(:,2); 
%plot(colX(~tfCircle),colY(~tfCircle),'xr');
plot(colX(~tfArea),colY(~tfArea),'xw'); % mark colonies that were discarded by area
text(50,100,num2str(colNum),'HorizontalAlignment','left','VerticalAlignment','bottom','BackgroundColor','w','fontsize',30);
strFileName = [strrep(strFileName,'_','-') '-QA.jpg'];
saveas(fh,strFileName)



