New update.

Enabled detecting a full map image.

Weed_Detector.py

@@ -1,108 +1,183 @@
-import os
-import cv2
-import zipfile
-import numpy as np
-import streamlit as st
-from io import BytesIO
-from PIL import Image
-from ultralytics import YOLO
-from utils import create_shapefile_with_latlon
-
-
-# Define paths
-path_to_store_bounding_boxes = 'detect/'
-path_to_save_shapefile = 'weed_detections.shp'
-
-# Ensure the output directories exist
-os.makedirs(path_to_store_bounding_boxes, exist_ok=True)
-
-# loading a custom model
-model = YOLO('new_yolov8_best.pt')
-
-# Mapping of class labels to readable names (assuming 'weeds' is class 1)
-class_names = ["citrus area", "trees", "weeds", "weeds and trees" ]
-
-
-# Streamlit UI
-st.title("Weed Detection and Shapefile Creation")
-
-# Input coordinates for image corners
-st.sidebar.header("Image Coordinates")
-top_left = st.sidebar.text_input("Top Left (lon, lat)", value="-48.8864783, -20.5906375")
-top_right = st.sidebar.text_input("Top Right (lon, lat)", value="-48.8855653, -20.5906264")
-bottom_right = st.sidebar.text_input("Bottom Right (lon, lat)", value="-48.8855534, -20.5914861")
-bottom_left = st.sidebar.text_input("Bottom Left (lon, lat)", value="-48.8864664, -20.5914973")
-
-# Convert input coordinates to tuples
-image_coords = [
-    tuple(map(float, top_left.split(','))),
-    tuple(map(float, top_right.split(','))),
-    tuple(map(float, bottom_right.split(','))),
-    tuple(map(float, bottom_left.split(',')))
-]
-
-# Upload image
-uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
-
-if uploaded_image is not None:
-    # Display uploaded image
-    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
-    img = Image.open(uploaded_image)
-    img_array = np.array(img)
-    image_height, image_width, _ = img_array.shape
-    temp_image_path = "temp_uploaded_image.png"
-    image = Image.open(uploaded_image)
-    image.save(temp_image_path)
-
-    # Perform weed detection on button click
-    if st.button("Detect Weeds"):
-        # Perform model prediction
-        results = model.predict(temp_image_path, imgsz=640, conf=0.2, iou=0.4)
-        results = results[0]
-        
-        weed_bboxes = []
-
-        for i, box in enumerate(results.boxes):
-            tensor = box.xyxy[0]
-            x1 = int(tensor[0].item())
-            y1 = int(tensor[1].item())
-            x2 = int(tensor[2].item())
-            y2 = int(tensor[3].item())
-            conf = box.conf[0].item()  # Confidence score
-            label = box.cls[0].item()  # Class label
-
-            # Debugging output to ensure boxes are detected
-            print(f"Box {i}: ({x1}, {y1}), ({x2}, {y2}), label: {label}, confidence: {conf}")
-
-            # Only process if the detected class is "weeds"
-            if class_names[int(label)] == "weeds":
-                print("weed detected")
-                # Draw bounding box on the image
-                cv2.rectangle(img_array, (x1, y1), (x2, y2), (255, 0, 255), 3)
-                # Save the bounding box coordinates
-                weed_bboxes.append((x1, y1, x2, y2))
-
-        # Save the image with bounding boxes
-        detected_image_path = os.path.join(path_to_store_bounding_boxes, "detected_image.png")
-        cv2.imwrite(detected_image_path, cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR))
-
-        # Display the image with bounding boxes
-        st.image(img_array, caption="Detected Weeds", use_column_width=True)
-
-        # Create shapefile with bounding boxes
-        create_shapefile_with_latlon(weed_bboxes, (image_width, image_height), image_coords, path_to_save_shapefile)
-
-        # Create ZIP file of the shapefile components
-        zip_buffer = BytesIO()
-        with zipfile.ZipFile(zip_buffer, 'w') as zip_file:
-            for filename in ['weed_detections.shp', 'weed_detections.shx', 'weed_detections.dbf']:
-                zip_file.write(filename, os.path.basename(filename))
-        zip_buffer.seek(0)
-
-        # Download ZIP file
-        st.download_button(
-            label="Download Shapefile ZIP",
-            data=zip_buffer,
-            file_name="weed_detections.zip",
-            mime="application/zip"
-        )
+import os
+import cv2
+import zipfile
+import numpy as np
+import streamlit as st
+from io import BytesIO
+from PIL import Image
+from ultralytics import YOLO
+from shapely.geometry import Polygon
+import shapefile
+import json
+import math
+from utils import create_shapefile_with_latlon
+
+# Increase the limit for PIL's decompression bomb protection
+Image.MAX_IMAGE_PIXELS = None
+
+# Define paths
+path_to_store_bounding_boxes = 'detect/'
+path_to_save_shapefile = 'weed_detections.shp'
+slice_folder = 'slices/'
+shapefile_folder = 'shapes/'
+
+# Ensure the output directories exist
+os.makedirs(path_to_store_bounding_boxes, exist_ok=True)
+os.makedirs(slice_folder, exist_ok=True)
+os.makedirs(shapefile_folder, exist_ok=True)
+
+# Loading a custom model
+model = YOLO('new_yolov8_best.pt')
+
+# Mapping of class labels to readable names (assuming 'weeds' is class 1)
+class_names = ["citrus area", "trees", "weeds", "weeds and trees"]
+
+# Streamlit UI
+st.title("Weed Detection and Shapefile Creation")
+
+# Input coordinates for image corners
+st.sidebar.header("Image Coordinates")
+top_left = st.sidebar.text_input("Top Left (lon, lat)", value="-48.8877415, -20.585013")
+top_right = st.sidebar.text_input("Top Right (lon, lat)", value="-48.8819718, -20.585013")
+bottom_right = st.sidebar.text_input("Bottom Right (lon, lat)", value="-48.8819718, -20.5968754")
+bottom_left = st.sidebar.text_input("Bottom Left (lon, lat)", value="-48.8877415, -20.5968754")
+
+# Convert input coordinates to tuples
+image_coords = [
+    tuple(map(float, top_left.split(','))),
+    tuple(map(float, top_right.split(','))),
+    tuple(map(float, bottom_right.split(','))),
+    tuple(map(float, bottom_left.split(',')))
+]
+
+# Upload image
+uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+
+def calculate_new_coordinates(original_coords, start_x, start_y, end_x, end_y, img_width, img_height):
+    lon_step = (original_coords[1][0] - original_coords[0][0]) / img_width
+    lat_step = (original_coords[0][1] - original_coords[3][1]) / img_height
+
+    new_top_left = (original_coords[0][0] + start_x * lon_step, original_coords[0][1] - start_y * lat_step)
+    new_top_right = (original_coords[0][0] + end_x * lon_step, original_coords[0][1] - start_y * lat_step)
+    new_bottom_right = (original_coords[0][0] + end_x * lon_step, original_coords[0][1] - end_y * lat_step)
+    new_bottom_left = (original_coords[0][0] + start_x * lon_step, original_coords[0][1] - end_y * lat_step)
+
+    return [new_top_left, new_top_right, new_bottom_right, new_bottom_left]
+
+def slice_image_and_coordinates(image_path, original_coords, slice_width=3000, slice_height=3000, output_folder='slices'):
+    os.makedirs(output_folder, exist_ok=True)
+    img = Image.open(image_path)
+    img_width, img_height = img.size
+
+    slice_coords = {}
+    slice_id = 0
+
+    num_slices_x = math.ceil(img_width / slice_width)
+    num_slices_y = math.ceil(img_height / slice_height)
+
+    for i in range(num_slices_y):
+        for j in range(num_slices_x):
+            start_x = j * slice_width
+            end_x = min(start_x + slice_width, img_width)
+            start_y = i * slice_height
+            end_y = min(start_y + slice_height, img_height)
+
+            box = (start_x, start_y, end_x, end_y)
+            cut_img = img.crop(box)
+
+            slice_filename = f'slice_{slice_id}.png'
+            cut_img.save(os.path.join(output_folder, slice_filename))
+
+            new_coords = calculate_new_coordinates(original_coords, start_x, start_y, end_x, end_y, img_width, img_height)
+            slice_coords[slice_filename] = new_coords
+
+            slice_id += 1
+
+    with open(os.path.join(output_folder, 'coordinates.json'), 'w') as json_file:
+        json.dump(slice_coords, json_file, indent=4)
+
+    return slice_coords
+
+def convert_pixel_to_latlon(x, y, image_width, image_height, image_coords):
+    top_left, top_right, bottom_right, bottom_left = image_coords
+
+    lon_top = top_left[0] + (top_right[0] - top_left[0]) * (x / image_width)
+    lon_bottom = bottom_left[0] + (bottom_right[0] - bottom_left[0]) * (x / image_width)
+    lat_left = top_left[1] + (bottom_left[1] - top_left[1]) * (y / image_height)
+    lat_right = top_right[1] + (bottom_right[1] - top_right[1]) * (y / image_height)
+
+    lon = lon_top + (lon_bottom - lon_top) * (y / image_height)
+    lat = lat_left + (lat_right - lat_left) * (x / image_width)
+
+    return lon, lat
+
+if uploaded_image is not None:
+    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
+    temp_image_path = "temp_uploaded_image.png"
+    image = Image.open(uploaded_image)
+    image.save(temp_image_path)
+
+    # Slice the image and save slices with their coordinates
+    slice_coords = slice_image_and_coordinates(temp_image_path, image_coords, slice_width=3000, slice_height=3000, output_folder=slice_folder)
+
+    if st.button("Detect Weeds"):
+        all_weed_bboxes = []
+
+        for slice_filename, coords in slice_coords.items():
+            slice_path = os.path.join(slice_folder, slice_filename)
+            image = cv2.imread(slice_path)
+            image_height, image_width, _ = image.shape
+
+            results = model.predict(slice_path, imgsz=640, conf=0.2, iou=0.4)
+            results = results[0]
+
+            weed_bboxes = []
+
+            for i, box in enumerate(results.boxes):
+                tensor = box.xyxy[0]
+                x1 = int(tensor[0].item())
+                y1 = int(tensor[1].item())
+                x2 = int(tensor[2].item())
+                y2 = int(tensor[3].item())
+                conf = box.conf[0].item()
+                label = box.cls[0].item()
+
+                if class_names[int(label)] == "weeds":
+                    cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 255), 3)
+                    weed_bboxes.append((x1, y1, x2, y2))
+
+            if weed_bboxes:
+                create_shapefile_with_latlon(weed_bboxes, (image_width, image_height), coords, f'shapes/{slice_filename.replace(".png", ".shp")}')
+                all_weed_bboxes.extend(weed_bboxes)
+
+            cv2.imwrite(os.path.join(path_to_store_bounding_boxes, slice_filename), image)
+
+        final_shapefile_path = path_to_save_shapefile
+        w = shapefile.Writer(final_shapefile_path)
+        w.field('id', 'C')
+
+        for slice_filename, coords in slice_coords.items():
+            shape_path = os.path.join(shapefile_folder, slice_filename.replace('.png', '.shp'))
+            if os.path.exists(shape_path):
+                r = shapefile.Reader(shape_path)
+                for shape_rec in r.iterShapeRecords():
+                    w.shape(shape_rec.shape)
+                    w.record(shape_rec.record[0])
+
+        w.close()
+
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, 'w') as zip_file:
+            for filename in ['weed_detections.shp', 'weed_detections.shx', 'weed_detections.dbf']:
+                zip_file.write(filename, os.path.basename(filename))
+        zip_buffer.seek(0)
+
+        st.download_button(
+            label="Download Shapefile ZIP",
+            data=zip_buffer,
+            file_name="weed_detections.zip",
+            mime="application/zip"
+        )
+
+        st.success("Weed detection completed and shapefile created successfully!")
+