Weed_Detector.py

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!")