% Experiment 10: Image segmentation using (a) Edge detection and (b) Thresholding
clear; close all; clc;

%% Part A: Using inbuilt MATLAB functions
% a. Edge detection
I = imread('circuit.tif');
BW_sobel   = edge(I, 'sobel');
BW_prewitt = edge(I, 'prewitt');
BW_roberts = edge(I, 'roberts');
BW_log     = edge(I, 'log');
[BW_canny, thresh_canny] = edge(I, 'canny');
fprintf('Canny thresholds chosen by edge(): low %.4f, high %.4f\n', thresh_canny(1), thresh_canny(2));

figure('Name', 'A_edges');
subplot(2, 3, 1); imshow(I);          title('Original (circuit.tif)');
subplot(2, 3, 2); imshow(BW_sobel);   title('Sobel');
subplot(2, 3, 3); imshow(BW_prewitt); title('Prewitt');
subplot(2, 3, 4); imshow(BW_roberts); title('Roberts');
subplot(2, 3, 5); imshow(BW_log);     title('Laplacian of Gaussian');
subplot(2, 3, 6); imshow(BW_canny);   title('Canny');

% b. Thresholding
C = imread('coins.png');
level = graythresh(C);                  % Otsu threshold in [0, 1]
BW_otsu = imbinarize(C, level);
G = imread('cameraman.tif');
levels = multithresh(G, 2);             % two thresholds -> three classes
seg3 = imquantize(G, levels);
fprintf('graythresh (Otsu) for coins.png = %.4f (%d gray levels)\n', level, round(level * 255));
fprintf('multithresh for cameraman.tif   = [%d %d]\n', levels(1), levels(2));

figure('Name', 'A_threshold');
subplot(2, 3, 1); imshow(C);       title('Original (coins.png)');
subplot(2, 3, 2); imhist(C);       title('Histogram');
subplot(2, 3, 3); imshow(BW_otsu); title(sprintf('Otsu (graythresh = %.3f)', level));
subplot(2, 3, 4); imshow(G);       title('Original (cameraman.tif)');
subplot(2, 3, 5); imhist(G);       title('Histogram');
subplot(2, 3, 6); imshow(label2rgb(seg3)); title('Multi-level (multithresh, 2 thresholds)');

%% Part B: Custom implementation
% a. Edge detection: Sobel gradient and Canny written from scratch
[sobelMag, sobelBW] = mySobel(I, 0.25);
cannyBW = edge_canny(I);

figure('Name', 'B_edges');
subplot(1, 4, 1); imshow(I);            title('Original');
subplot(1, 4, 2); imshow(sobelMag, []); title('Custom Sobel gradient magnitude');
subplot(1, 4, 3); imshow(sobelBW);      title('Custom Sobel edges');
subplot(1, 4, 4); imshow(cannyBW);      title('Custom Canny');

% b. Thresholding: basic global (iterative) and Otsu, without inbuilt functions
[Tg, iters] = basicGlobalThreshold(C, 0.5);
[kStar, eta, sigmaB] = otsuThreshold(C);
BW_global = C > Tg;
BW_myOtsu = C > kStar;
[k1, k2] = otsuTwoLevel(G);
seg3B = 1 + (G > k1) + (G > k2);
fprintf('\nBasic global threshold : T = %.2f after %d iterations\n', Tg, iters);
fprintf('Custom Otsu threshold  : k* = %d, separability = %.4f\n', kStar, eta);
fprintf('Inbuilt Otsu threshold : %d (graythresh x 255)\n', round(level * 255));
fprintf('Custom two-level Otsu  : [%d %d]\n', k1, k2);
fprintf('Pixels that differ, custom Otsu vs imbinarize: %d\n', nnz(BW_myOtsu ~= BW_otsu));

figure('Name', 'B_threshold');
subplot(2, 3, 1); imshow(BW_global); title(sprintf('Basic global, T = %.1f', Tg));
subplot(2, 3, 2); imshow(BW_myOtsu); title(sprintf('Custom Otsu, k* = %d', kStar));
subplot(2, 3, 3); plot(0:255, sigmaB, 'b', 'LineWidth', 1.3); hold on;
xline(kStar, 'r--', sprintf('k* = %d', kStar)); hold off; xlim([0 255]);
title('Between-class variance \sigma_B^2(k)'); xlabel('k');
subplot(2, 3, 4); imshow(G); title('Original (cameraman.tif)');
subplot(2, 3, 5); imshow(label2rgb(seg3B)); title(sprintf('Custom two-level Otsu [%d %d]', k1, k2));
subplot(2, 3, 6); bar(0:255, myHist(G), 1, 'k'); hold on;
xline(k1, 'r--'); xline(k2, 'r--'); hold off; xlim([0 255]); title('Histogram with thresholds');

%% Local functions
function h = myHist(img)
    h = accumarray(double(img(:)) + 1, 1, [256 1]);
end

function [mag, BW] = mySobel(img, frac)
    img = double(img);
    Kx = [-1 0 1; -2 0 2; -1 0 1];
    Ky = [-1 -2 -1; 0 0 0; 1 2 1];
    P = padarray(img, [1 1], 'replicate');
    gx = zeros(size(img)); gy = zeros(size(img));
    for x = 1:size(img, 1)
        for y = 1:size(img, 2)
            region = P(x:x+2, y:y+2);
            gx(x, y) = sum(region .* Kx, 'all');
            gy(x, y) = sum(region .* Ky, 'all');
        end
    end
    mag = sqrt(gx.^2 + gy.^2);
    BW = mag > frac * max(mag(:));
end

function BW = edge_canny(I)
    % smoothing, Sobel gradient, non-maximum suppression, hysteresis thresholding
    I_double = double(I);
    g_filter = [1 2 1; 2 4 2; 1 2 1] / 16;
    I_smooth = conv2(I_double, g_filter, 'same');
    Kx = [-1 0 1; -2 0 2; -1 0 1];
    Ky = [1 2 1; 0 0 0; -1 -2 -1];
    Ix = conv2(I_smooth, Kx, 'same');
    Iy = conv2(I_smooth, Ky, 'same');
    Mag = hypot(Ix, Iy);
    Ang = atan2d(Iy, Ix);
    Ang(Ang < 0) = Ang(Ang < 0) + 180;
    [rows, cols] = size(Mag);
    NMS = zeros(rows, cols);
    for r = 2:rows-1
        for c = 2:cols-1
            a = Ang(r, c);
            if (a >= 0 && a < 22.5) || (a >= 157.5 && a <= 180)
                v1 = Mag(r, c+1); v2 = Mag(r, c-1);
            elseif (a >= 22.5 && a < 67.5)
                v1 = Mag(r-1, c+1); v2 = Mag(r+1, c-1);
            elseif (a >= 67.5 && a < 112.5)
                v1 = Mag(r+1, c); v2 = Mag(r-1, c);
            else
                v1 = Mag(r-1, c-1); v2 = Mag(r+1, c+1);
            end
            if Mag(r, c) >= v1 && Mag(r, c) >= v2
                NMS(r, c) = Mag(r, c);
            end
        end
    end
    high_th = max(NMS(:)) * 0.15;
    low_th = high_th * 0.4;
    strong = NMS >= high_th;
    weak = (NMS >= low_th) & (NMS < high_th);
    BW = strong;
    [r_idx, c_idx] = find(strong);
    stack = [r_idx, c_idx];
    while ~isempty(stack)
        curr = stack(end, :);
        stack(end, :) = [];
        for dr = -1:1
            for dc = -1:1
                nr = curr(1) + dr;
                nc = curr(2) + dc;
                if nr > 0 && nr <= rows && nc > 0 && nc <= cols && weak(nr, nc) && ~BW(nr, nc)
                    BW(nr, nc) = true;
                    stack = [stack; nr, nc]; %#ok<AGROW>
                end
            end
        end
    end
end

function [T, iters] = basicGlobalThreshold(img, dT)
    img = double(img(:));
    T = mean(img);                      % initial estimate
    iters = 0;
    while true
        iters = iters + 1;
        m1 = mean(img(img > T));
        m2 = mean(img(img <= T));
        Tnew = 0.5 * (m1 + m2);
        if abs(Tnew - T) < dT
            T = Tnew;
            break;
        end
        T = Tnew;
    end
end

function [kStar, eta, sigmaB] = otsuThreshold(img)
    p = myHist(img) / numel(img);       % normalized histogram
    i = (0:255)';
    P1 = cumsum(p);                     % cumulative sums
    mk = cumsum(i .* p);                % cumulative means
    mG = mk(end);                       % global mean
    sigmaB = (mG * P1 - mk).^2 ./ (P1 .* (1 - P1));
    sigmaB(~isfinite(sigmaB)) = 0;
    best = find(sigmaB == max(sigmaB));
    kStar = round(mean(best)) - 1;      % average if the maximum is not unique
    sigmaG = sum((i - mG).^2 .* p);
    eta = max(sigmaB) / sigmaG;         % separability measure
end

function [k1, k2] = otsuTwoLevel(img)
    % exhaustive search for the two thresholds that maximize sigma_B^2
    p = myHist(img) / numel(img);
    i = (0:255)';
    P = cumsum(p); S = cumsum(i .* p);
    mG = S(end);
    best = -1; k1 = 0; k2 = 0;
    for a = 1:254
        for b = a+1:255
            w = [P(a), P(b) - P(a), 1 - P(b)];
            s = [S(a), S(b) - S(a), mG - S(b)];
            if all(w > 0)
                sb = sum(w .* (s ./ w - mG).^2);
                if sb > best
                    best = sb; k1 = a - 1; k2 = b - 1;
                end
            end
        end
    end
end