Update app.py

app.py

@@ -107,14 +107,33 @@ with row3_col5:
 # Output and Interaction
 ###
 #####
-if st.session_state.image is not None:
+# Function to break the image into tiles
+def tile_image(image, tile_size, overlap):
+    # Get the image size
+    width, height = image.size
+    
+    # Create a list to store the tiles
+    tiles = []
+    
+    # Loop through the image and break it into tiles
+    for x in range(0, width, tile_size - overlap):
+        for y in range(0, height, tile_size - overlap):
+            # Get the tile
+            tile = image.crop((x, y, x + tile_size, y + tile_size))
+            tiles.append(tile)
+    
+    return tiles
         
+if st.session_state.image is not None:
+    # tile images
+    tiles = tile_image(image=st.session_state.image, tile_size=640, overlap=40)
+    
     st.write("\n\n")
     components.html("""<hr style="height:2px;border:none;color:#333;background-color:#333;" /> """)
 
     # model
     model = YOLO("/home/user/app/models/head_model/best.pt")  # load
-    results = model.predict([st.session_state.image]) # predict
+    results = model.predict(tiles) # predict
     
     
     _, col_1, col_2, _ = st.columns([1, 4, 4, 1], gap="medium")