%% calcColoniesOnPlate.m
% segments colonies on an agar plate using a watershed algorithm (count blue and white
% colonies)
% Amir Mitchell 07/08/2019

%% User definitions
strFileName = '20210307AM_MB_0007_12_Tumor12_BlueWhite_Backgrnd2_100000Dil.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;
pixelCutoff = nan; % pixel intensity to classify blue/white (keep at nan, unless script fails) 

%% open image, convert to gray, clean up, convert to BW
original = imresize(imread(strFileName),0.5); % open file
blue = original(:,:,3); % blue channel for segmeting all colonies (blue/white colonies both look white)

%% crop the image by user defined bounding box
figure; imshow(blue); 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(blue,[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');

%% we also make a crop version of the red channel (red allows classifying blue/whire colonies)
original_crop = imcrop(original(:,:,:),[x y l l]);
red = original_crop(:,:,1);

%% 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));
trueCol = [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

%% 
medianRedSignal = zeros(size(cc));
for i=1:length(cc)
    medianRedSignal(i) = median(red(cc(i).PixelIdxList));
end

% decideo on the cutoff intensity to classify blue/white
if(isnan(pixelCutoff))
    pixelCutoff = mean([min(medianRedSignal(trueCol)) max(medianRedSignal(trueCol))]);
end

inxWhite = [medianRedSignal>pixelCutoff & trueCol];
inxBlue = [medianRedSignal<=pixelCutoff & trueCol];

nWhite = sum(inxWhite);
nBlue = sum(inxBlue);

figure; hold on;
histogram(medianRedSignal(trueCol),100);
plot([pixelCutoff pixelCutoff],[-5 1],'r');
text(mean(medianRedSignal(inxWhite)),-2,['N blue = ' num2str(nWhite)]);
text(mean(medianRedSignal(inxBlue)),-2,['N blue = ' num2str(nBlue)]);

%% make a QA figure (only detected colonies are colored)
colNum = sum(trueCol);
fh = figure;
pause(1);
subplot(1,2,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);

subplot(1,2,2);

imshow(original_crop); 
hold on;
xy = cat(1,cc.Centroid)
plot(xy(inxWhite,1),xy(inxWhite,2),'ow','markersize',10);
plot(xy(inxBlue,1),xy(inxBlue,2),'ob','markersize',10);
text(50,100,[num2str(nBlue) '/' num2str(nWhite)] ,'HorizontalAlignment','left','VerticalAlignment','bottom','BackgroundColor','w','fontsize',30);


strTitle = [strrep(strFileName,'_','-') ' (blue=' num2str(nBlue) ', white=' num2str(nWhite) ')'];
title(strTitle); % title by file name, replace '_' with '-'


set(gcf,'Position',[200 200 1100 600])
strFileName = [strrep(strFileName,'_','-') '-QA.jpg'];
saveas(fh,strFileName)



