Upload Colony_Analyzer_AI2_HF.py

Colony_Analyzer_AI2_HF.py

@@ -8,37 +8,90 @@ Created on Thu Mar 20 14:23:27 2025
 
 import os
 import cv2
+from PIL import Image
+
+
+def ensure_minimum_size(img, min_width=512, min_height=512):
+    width, height = img.size
+    if width < min_width or height < min_height:
+        return img.resize((max(width, min_width), max(height, min_height)), Image.BICUBIC)
+    return img
+
+def pad(img_np):
+    """
+    Pads a numpy image (grayscale or RGB) to 2048x1536 (width x height) with white pixels.
+    Pads at the bottom and right as needed.
+    """
+    target_width, target_height = 2048, 1536
+    height, width = img_np.shape[:2]
+    pad_bottom = max(0, target_height - height)
+    pad_right = max(0, target_width - width)
+    # Padding: (top, bottom, left, right)
+    if img_np.ndim == 3:
+        # Color image (H, W, 3)
+        border_value = [255, 255, 255]
+    else:
+        # Grayscale image (H, W)
+        border_value = 255
+
+    padded = cv2.copyMakeBorder(
+        img_np, 
+        top=0, bottom=pad_bottom,
+        left=0, right=pad_right,
+        borderType=cv2.BORDER_CONSTANT,
+        value=border_value
+    )
+    return padded
+
 
 #this is the huggingface version
-def cut_img(img):
+def cut_img(img, patch_size=512):
     img_map = {}
     width, height = img.size
-    i_num = height // 512
-    j_num = width // 512
+    i_num = height // patch_size
+    j_num = width // patch_size
     count = 1
     for i in range(i_num):
         for j in range(j_num):
-            cropped_img = img.crop((512*j, 512*i, 512*(j+1), 512*(i+1)))
+            cropped_img = img.crop((
+                patch_size * j,
+                patch_size * i,
+                patch_size * (j + 1),
+                patch_size * (i + 1)
+            ))
             img_map[count] = cropped_img
-            #print(type(cropped_img))
             count += 1
-    return img_map
+    return img_map, i_num, j_num  # Return rows and cols for stitching
     
 import numpy as np
 
-def stitch(img_map):
-    rows = [
-        np.hstack([img_map[1], img_map[2], img_map[3], img_map[4]]),  # First row (images 0 to 3)
-        np.hstack([img_map[5], img_map[6], img_map[7], img_map[8]]),  # Second row (images 4 to 7)
-        np.hstack([img_map[9], img_map[10], img_map[11], img_map[12]])  # Third row (images 8 to 11)
-    ]
-    # Stack rows vertically
-    return(np.vstack(rows))
-
-
-
+import numpy as np
 from PIL import Image
 
+def stitch(img_map, i_num, j_num, min_width=2048, min_height=1536):
+    tiles = []
+    count = 1
+    for i in range(i_num):
+        row_tiles = []
+        for j in range(j_num):
+            tile = np.array(img_map[count])
+            row_tiles.append(tile)
+            count += 1
+        row_img = np.hstack(row_tiles)
+        tiles.append(row_img)
+    stitched = np.vstack(tiles)
+
+    # Pad the stitched image if it's less than min_width/min_height
+    h, w = stitched.shape[:2]
+    pad_h = max(0, min_height - h)
+    pad_w = max(0, min_width - w)
+    if pad_h > 0 or pad_w > 0:
+        # Pad as (top, bottom), (left, right), (channels)
+        if stitched.ndim == 3:
+            stitched = np.pad(stitched, ((0, pad_h), (0, pad_w), (0, 0)), 'constant')
+        else:
+            stitched = np.pad(stitched, ((0, pad_h), (0, pad_w)), 'constant')
+    return stitched
 
 
 import matplotlib.pyplot as plt
@@ -214,21 +267,24 @@ def main(args):
     min_circ = args[2]
     do_necrosis = args[3]
     colonies = {}
-    img_map = cut_img(args[0])
+    img_map, i_num, j_num = cut_img(ensure_minimum_size(args[0], min_width=512, min_height=512))
     for z in img_map:
         img_map[z] = eval_img(img_map[z])
     del z
-    p = stitch(img_map)
+    p = stitch(img_map, i_num, j_num)
     colonies = analyze_colonies(p, min_size, min_circ, np.array(args[0]))
     if len(colonies) <=0:
+    	img = np.array(pad(args[0]))
     	caption = np.ones((150, 2048, 3), dtype=np.uint8) * 255  # Multiply by 255 to make it white
     	cv2.putText(caption, 'No colonies detected.', (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 3)
-    	cv2.imwrite('results.png', np.vstack((np.array(args[0]), caption)))
+    	cv2.imwrite('results.png', np.vstack((img, caption)))
     	colonies = pd.DataFrame({"Colony Number":[], 'Colony volume':[], "colony_area":[],'mean_pixel_value':[], "centroid":[], "necrotic_area":[],"percent_necrotic":[]})
     	with pd.ExcelWriter('results.xlsx') as writer:
         	colonies.to_excel(writer, sheet_name="Colony data", index=False)
-    	return(np.vstack((args[0], caption)), 'results.png', 'results.xlsx')
-    img = np.array(args[0])
+    	return(np.vstack((img, caption)), 'results.png', 'results.xlsx')
+
+    img = np.array(pad(args[0]))
+
     img = cv2.copyMakeBorder(img,top=0, bottom=10,left=0,right=10, borderType=cv2.BORDER_CONSTANT,  value=[255, 255, 255]) 
     #print(colonies.to_string())