import numpy as np


def correlation(image, angrange, brange, sigma):
    xsize, ysize = image.shape
    
    subimage = image.astype(np.float32)
    
    medianint = np.median(image[image > 0])
    angles = np.arange(angrange[0], angrange[2], angrange[1])
    trans = np.arange(brange[0], brange[2], brange[1])
    corrmap = np.zeros((len(angles), len(trans)))
    
    for thetaind in range(len(angles)):
        for bind in range(len(trans)):
            theta = angles[thetaind]
            b = trans[bind]
            refline = reference_line_gen(xsize, ysize, theta, b, sigma)
            temp = refline * subimage
            temp2 = np.sum(temp)
            temp3 = temp2 / medianint
            corrmap[thetaind, bind] = temp3
            
    return corrmap


def reference_line_gen(xsize, ysize, theta, b, sigma):
    Xmesh, Ymesh = np.meshgrid(np.arange(1, xsize+1), np.arange(1, ysize+1))
    Xmesh = Xmesh - xsize / 2
    Ymesh = Ymesh - ysize / 2
    
    angle = np.deg2rad(theta)
    dist = b - (Xmesh * np.sin(angle)) - (Ymesh * np.cos(angle))
    distance = dist ** 2 / (sigma ** 2)
    result = np.exp(-distance)
    
    return result


def sub_window_analysis(image, angrange, brange, sigma):
    imgsize = np.shape(image)
    sgimgsize = np.array(imgsize, dtype=np.float32)
    sgsubimgsz = np.array(image, dtype=np.float32)
    sgtestimage = np.array(image, dtype=np.float32)
    sigma = np.float32(sigma)

    if imgsize[0] % subimgsz[0] != 0 or imgsize[1] % subimgsz[1] != 0:
        print("The subimage size does not match the whole image size")
        return

    angles = np.arange(angrange[0], angrange[1], angrange[2])
    trans = np.arange(brange[0], brange[1], brange[2])
    corr = np.zeros((len(angles), len(trans), imgsize[0]//subimgsz[0], imgsize[1]//subimgsz[1]))
    xsize = sgimgsize[1]
    ysize = sgimgsize[0]
    Xmesh, Ymesh = np.meshgrid(np.arange(0,xsize), np.arange(0,ysize))
    Xmesh = np.mod(Xmesh, subimgsz[1])
    Ymesh = np.mod(Ymesh, subimgsz[0])
    Xmesh = Xmesh - sgsubimgsz[1]/2
    Ymesh = Ymesh - sgsubimgsz[0]/2

    for thetaind in range(len(angles)):
        for bind in range(len(trans)):
            theta = angles[thetaind]
            b = np.float32(trans[bind])
            angle = (theta * np.pi) / 180
            dist = b - (Xmesh * np.sin(angle)) - (Ymesh * np.cos(angle))
            distpw = dist ** 2 / (sigma ** 2)
            refline = np.exp(-distpw)
            temp = refline * sgtestimage
            temp2 = np.reshape(temp, (subimgsz[0], imgsize[0]//subimgsz[0], subimgsz[1], imgsize[1]//subimgsz[1]))
            temp3 = np.sum(temp2, axis=1)
            temp4 = np.squeeze(np.sum(temp3, axis=2))
            corr[thetaind, bind, :, :] = temp4

    corrsz = np.shape(corr)
    spcorr = np.reshape(np.transpose(corr, (0, 2, 1, 3)), (corrsz[0] * corrsz[2], corrsz[1] * corrsz[3]))
    colsubimgs = np.reshape(sgtestimage, (subimgsz[0], imgsize[0] // subimgsz[0], subimgsz[1], imgsize[1] // subimgsz[1]))
    
    return corrsz, spcorr, colsubimgs

# Example Usage
image = cv2.imread("0.0/01.pgm", cv2.COLOR_BGR2GRAY)
angrange = (0, 10, 180)   # Range of angles for correlation analysis
brange = (0, 1, 10)       # Range of translation values for correlation analysis
sigma = 1.0               # Sigma value for reference line generation

# Perform whole image correlation analysis
corr_map = correlation(image, angrange, brange, sigma)
print("Correlation map shape:", corr_map.shape)

# Extract maximum intensity peaks from the correlation map
label_map = max_intensity_finding(corr_map)
print("Label map shape:", label_map.shape)

# Plot the correlation map and label map
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.imshow(corr_map, cmap="gray")
plt.title("Correlation Map")
plt.colorbar()
plt.subplot(1, 2, 2)
plt.imshow(label_map, cmap="gray")
plt.title("Label Map")
plt.colorbar()
plt.show()

# Perform sub-window analysis
subimgsz = (10, 10)  # Sub-window size for analysis
corrsz, spcorr, colsubimgs = sub_window_analysis(image, angrange, brange, sigma)
print("Spaced correlation shape:", spcorr.shape)
print("Column sub-images shape:", colsubimgs.shape)

ПОЛУЧАЮ ОШИБКУ:
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-21-81e00d182b41> in <cell line: 110>()
    108 # Perform sub-window analysis
    109 subimgsz = (10, 10)  # Sub-window size for analysis
--> 110 corrsz, spcorr, colsubimgs = sub_window_analysis(image, angrange, brange, sigma)
    111 print("Spaced correlation shape:", spcorr.shape)
    112 print("Column sub-images shape:", colsubimgs.shape)

<ipython-input-21-81e00d182b41> in sub_window_analysis(image, angrange, brange, sigma)
     72             temp3 = np.sum(temp2, axis=1)
     73             temp4 = np.squeeze(np.sum(temp3, axis=2))
---> 74             corr[thetaind, bind, :, :] = temp4
     75 
     76     corrsz = np.shape(corr)

ValueError: could not broadcast input array from shape (10,10) into shape (50,50)
ИСПРАВЬ КОД