Vedika8/image-processing-pipeline

1import sys
2import os
3sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
4
5import cv2
6import torch
7from PIL import Image
8import numpy as np
9from ultralytics import YOLO
10from models.segmentation_model import SegmentationModel
11
12def test_segmentation():
13    # Initialize the segmentation model
14    model = SegmentationModel()
15
16    # Specify the path to your test image
17    image_path = 'E:\downloads2\waser 5\waser\data\input_images\dog-park-petting-dog.jpg'
18
19    # Run segmentation
20    masks, boxes, labels, class_names = model.segment_image(image_path)
21
22    # Output the results
23    print(f"Detected Classes: {class_names}")
24    print(f"Bounding Boxes: {boxes}")
25    print(f"Labels: {labels}")
26    print(f"Masks Shape: {masks.shape if masks.size else 'No masks detected'}")
27
28    # Optionally, visualize the results
29    if boxes.size > 0:
30        image = cv2.imread(image_path)
31        for i, box in enumerate(boxes):
32            x1, y1, x2, y2 = map(int, box)
33            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
34            cv2.putText(image, class_names[i], (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
35        cv2.imshow("Segmented Image", image)
36        cv2.waitKey(0)
37        cv2.destroyAllWindows()
38
39if __name__ == "__main__":
40    test_segmentation()