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