Team Ai
Apppublic

todap/Image-Processing-Pipeline

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes
visualization.py54 linesDownload Raw Back to utils
1import matplotlib.pyplot as plt2import pandas as pd3import cv24from ultralytics import YOLO5from PIL import Image6def visualize_detections(image_path, output_path):7    8    model = YOLO('yolov8s.pt')  # You can change this to other YOLOv8 models as needed9    # Read the image10    image = cv2.imread(image_path)11 12    # Run YOLOv8 inference on the image13    results = model(image)14 15    # Process the results and draw bounding boxes16    for result in results:17        boxes = result.boxes.cpu().numpy()18        for box in boxes:19            x1, y1, x2, y2 = map(int, box.xyxy[0])20            confidence = float(box.conf[0])21            class_id = int(box.cls[0])22            class_name = model.names[class_id]23 24            # Draw bounding box25            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)26 27            # Prepare label28            label = f"{class_name}"29            30            # Get label size31            (label_width, label_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)32 33            # Draw filled rectangle for label background34            cv2.rectangle(image, (x1, y1 - label_height - baseline), (x1 + label_width, y1), (0, 255, 0), cv2.FILLED)35 36            # Put label text37            cv2.putText(image, label, (x1, y1 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)38 39    # Save the output image40    cv2.imwrite(output_path, image)41 42def visualize_segmentation(image, masks, output_file):43    #plt.imshow(image)44    for mask in masks:45        plt.imshow(mask, alpha=0.5)46    plt.axis('off')47    plt.savefig(output_file,bbox_inches='tight', pad_inches=0)48    plt.close()49 50 51def create_summary_table(mapped_data, output_file):52    df = pd.DataFrame.from_dict(mapped_data, orient='index')53    df.to_csv(output_file)54