Spaces:
Sleeping
Sleeping
| 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): | |
| boxes = [] | |
| with open(file_path, 'r') as f: | |
| for line in f: | |
| parts = line.strip().split() | |
| class_name = COCO_CLASSES[int(parts[0])] | |
| x, y, w, h = map(float, parts[1:5]) | |
| boxes.append((class_name, x, y, w, h)) | |
| return boxes | |
| 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 | |
| COCO_CLASSES = [ | |
| 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', | |
| 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', | |
| 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', | |
| 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', | |
| 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', | |
| 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', | |
| 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', | |
| 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', | |
| 'hair drier', 'toothbrush' | |
| ] | |
| def read_yolo_boxes(file_path): | |
| boxes = [] | |
| with open(file_path, 'r') as f: | |
| for line in f: | |
| parts = line.strip().split() | |
| class_name = COCO_CLASSES[int(parts[0])] | |
| x, y, w, h = map(float, parts[1:5]) | |
| boxes.append((class_name, x, y, w, h)) | |
| return boxes | |
| def main(): | |
| st.title("YOLO Analysis App") | |
| uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png"]) | |
| if uploaded_file is not None: | |
| # Read the image as BGR | |
| image = cv2.imdecode(np.fromstring(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR) | |
| # Convert BGR to RGB | |
| image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| st.image(image_rgb, 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_rgb, yolo_boxes, segment, img_width, img_height, threshold) | |
| st.pyplot(fig) | |
| st.subheader("Analysis Results:") | |
| for class_name, x, y, w, h in yolo_boxes: | |
| result = box_segment_relationship((0, x, y, w, h), segment, img_width, img_height, threshold) | |
| st.write(f"{class_name} at ({x:.2f}, {y:.2f}) 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() |