Spaces:
Sleeping
Sleeping
File size: 1,781 Bytes
26932f6 |
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 |
import matplotlib.pyplot as plt
import pandas as pd
import cv2
from ultralytics import YOLO
from PIL import Image
def visualize_detections(image_path, output_path):
model = YOLO('yolov8s.pt') # You can change this to other YOLOv8 models as needed
# Read the image
image = cv2.imread(image_path)
# Run YOLOv8 inference on the image
results = model(image)
# Process the results and draw bounding boxes
for result in results:
boxes = result.boxes.cpu().numpy()
for box in boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
confidence = float(box.conf[0])
class_id = int(box.cls[0])
class_name = model.names[class_id]
# Draw bounding box
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
# Prepare label
label = f"{class_name}"
# Get label size
(label_width, label_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
# Draw filled rectangle for label background
cv2.rectangle(image, (x1, y1 - label_height - baseline), (x1 + label_width, y1), (0, 255, 0), cv2.FILLED)
# Put label text
cv2.putText(image, label, (x1, y1 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
# Save the output image
cv2.imwrite(output_path, image)
def visualize_segmentation(image, masks, output_file):
#plt.imshow(image)
for mask in masks:
plt.imshow(mask, alpha=0.5)
plt.axis('off')
plt.savefig(output_file,bbox_inches='tight', pad_inches=0)
plt.close()
def create_summary_table(mapped_data, output_file):
df = pd.DataFrame.from_dict(mapped_data, orient='index')
df.to_csv(output_file)
|