Spaces:
Sleeping
Sleeping
File size: 6,175 Bytes
858c475 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import os
import cv2
import numpy as np
import streamlit as st
import matplotlib.pyplot as plt
from shapely.geometry import Polygon, box as shapely_box
import subprocess
# ... (previous functions remain unchanged)
def extract_class_0_coordinates(filename):
class_0_coordinates = []
current_class = None
with open(filename, 'r') as file:
for line in file:
parts = line.strip().split()
if len(parts) == 0:
continue
if parts[0] == '0':
coordinates = [float(x) for x in parts[1:]]
class_0_coordinates.extend(coordinates)
return class_0_coordinates
def run_yolo_models1(img):
# Run YOLOv9 segmentation
os.system(f"python segment/predict.py --source {img} --img 640 --device cpu --weights models/segment/best-2.pt --name yolov9_c_640_detect --exist-ok --save-txt")
# Run YOLOv9 detection
os.system(f"python detect.py --source {img} --img 640 --device cpu --weights models/detect/yolov9-s-converted.pt --name yolov9_c_640_detect --exist-ok --save-txt")
def parse_yolo_box(box_string):
"""Parse a YOLO format bounding box string."""
values = list(map(float, box_string.split()))
if len(values) < 5:
raise ValueError(f"Expected at least 5 values, got {len(values)}")
return values[0], values[1], values[2], values[3], values[4]
def read_yolo_boxes(file_path):
"""Read YOLO format bounding boxes from a file."""
with open(file_path, 'r') as f:
return [parse_yolo_box(line.strip()) for line in f if line.strip()]
def yolo_to_pixel_coord(x, y, img_width, img_height):
"""Convert a single YOLO coordinate to pixel coordinate."""
return int(x * img_width), int(y * img_height)
def yolo_to_pixel_coords(x_center, y_center, width, height, img_width, img_height):
"""Convert YOLO format coordinates to pixel coordinates."""
x1 = int((x_center - width / 2) * img_width)
y1 = int((y_center - height / 2) * img_height)
x2 = int((x_center + width / 2) * img_width)
y2 = int((y_center + height / 2) * img_height)
return x1, y1, x2, y2
def box_segment_relationship(yolo_box, segment, img_width, img_height, threshold):
"""Check the relationship between a bounding box and a segmented area."""
class_id, x_center, y_center, width, height = yolo_box
x1, y1, x2, y2 = yolo_to_pixel_coords(x_center, y_center, width, height, img_width, img_height)
pixel_segment = convert_segment_to_pixel(segment, img_width, img_height)
segment_polygon = Polygon(zip(pixel_segment[::2], pixel_segment[1::2]))
box_polygon = shapely_box(x1, y1, x2, y2)
if box_polygon.intersects(segment_polygon):
return "intersecting"
elif box_polygon.distance(segment_polygon) <= threshold:
return "obstructed"
else:
return "not touching"
def convert_segment_to_pixel(segment, img_width, img_height):
"""Convert segment coordinates from YOLO format to pixel coordinates."""
pixel_segment = []
for i in range(0, len(segment), 2):
x, y = yolo_to_pixel_coord(segment[i], segment[i+1], img_width, img_height)
pixel_segment.extend([x, y])
return pixel_segment
def plot_boxes_and_segment(image, yolo_boxes, segment, img_width, img_height, threshold):
"""Plot the image with intersecting boxes, obstructed boxes, and segment."""
fig, ax = plt.subplots(figsize=(12, 8))
ax.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
pixel_segment = convert_segment_to_pixel(segment, img_width, img_height)
ax.plot(pixel_segment[::2] + [pixel_segment[0]], pixel_segment[1::2] + [pixel_segment[1]], 'g-', linewidth=2, label='Rail Zone')
colors = {'intersecting': 'r', 'obstructed': 'y', 'not touching': 'b'}
labels = {'intersecting': 'Intersecting Box', 'obstructed': 'Obstructed Box', 'not touching': 'Non-interacting Box'}
for yolo_box in yolo_boxes:
class_id, x_center, y_center, width, height = yolo_box
x1, y1, x2, y2 = yolo_to_pixel_coords(x_center, y_center, width, height, img_width, img_height)
relationship = box_segment_relationship(yolo_box, segment, img_width, img_height, threshold)
color = colors[relationship]
label = labels[relationship]
ax.add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, fill=False, edgecolor=color, linewidth=2, label=label))
ax.legend()
ax.axis('off')
plt.tight_layout()
return fig
def main():
st.title("YOLO Analysis App")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png"])
if uploaded_file is not None:
image = cv2.imdecode(np.fromstring(uploaded_file.read(), np.uint8), 1)
st.image(image, caption='Uploaded Image', use_column_width=True)
if st.button('Run Analysis'):
with st.spinner("Running detection..."):
img_height, img_width = image.shape[:2]
# Save the uploaded image temporarily
temp_image_path = "temp_image.jpg"
cv2.imwrite(temp_image_path, image)
# Run YOLO models
run_yolo_models1(temp_image_path)
label_path = 'runs/predict-seg/yolov9_c_640_detect/labels/temp_image.txt'
label_path2 = 'runs/detect/yolov9_c_640_detect/labels/temp_image.txt'
segment = extract_class_0_coordinates(label_path)
yolo_boxes = read_yolo_boxes(label_path2)
threshold = 10 # Set threshold (in pixels)
fig = plot_boxes_and_segment(image, yolo_boxes, segment, img_width, img_height, threshold)
st.pyplot(fig)
st.subheader("Analysis Results:")
for yolo_box in yolo_boxes:
result = box_segment_relationship(yolo_box, segment, img_width, img_height, threshold)
st.write(f"Box {yolo_box} is {result} the segment.")
# Clean up temporary files
os.remove(temp_image_path)
os.remove(label_path)
os.remove(label_path2)
if __name__ == "__main__":
main() |