Skip to content
Snippets Groups Projects
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
Q2main.py 6.82 KiB
import cv2
import matplotlib.pyplot as plt
import numpy as np
from skimage.color import rgb2gray
from skimage.feature import blob_log

class ImageProcessor:
    def __init__(self, image_path):
        self.image = cv2.imread(image_path)
        self.image_rgb = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
        self.image_hsv = cv2.cvtColor(self.image, cv2.COLOR_BGR2HSV)

    def apply_median_filter(self, kernel_size = 5):
        filtered_image = cv2.medianBlur(self.image_rgb,kernel_size)
        return filtered_image

    def get_color_channels(self):
        red_channel = self.image_rgb[:, :, 0]
        green_channel = self.image_rgb[:, :, 1]
        blue_channel = self.image_rgb[:, :, 2]
        return red_channel, green_channel, blue_channel
    
    def convert_to_hsv(self):
        return cv2.cvtColor(self.image, cv2.COLOR_BGR2HSV)
        

    def calculate_snr(self, channel):
        mean = np.mean(channel)
        stddev = np.std(channel)
        snr = mean / stddev
        return snr 
    
    def get_color_masks(self):
        lower_red1 = np.array([0, 100, 100])
        upper_red1 = np.array([10, 255, 255])
        lower_red2 = np.array([170, 100, 100])
        upper_red2 = np.array([180, 255, 255])

        lower_yellow = np.array([22, 100, 100])
        upper_yellow = np.array([32, 255, 255])

        lower_green = np.array([45, 100, 100])
        upper_green = np.array([72, 255, 255])

        red_mask = cv2.inRange(self.image_hsv, lower_red1, upper_red1) | cv2.inRange(self.image_hsv, lower_red2, upper_red2)
        yellow_mask =cv2.inRange(self.image_hsv, lower_yellow, upper_yellow)
        green_mask = cv2.inRange(self.image_hsv, lower_green, upper_green)

        return red_mask, yellow_mask, green_mask

def detect_blobs(mask, min_sigma=5, max_sigma=50, num_sigma = 10, threshold=0.5):
    mask = mask.astype(np.float64) / 255.0 
    #print('mask made')
    blobs = blob_log(mask, min_sigma=min_sigma, max_sigma=max_sigma, num_sigma=num_sigma, threshold=threshold)
    #print('blobs detected')
    if len(blobs) > 0:
        blobs[:,2] = blobs[:, 2]*np.sqrt(2)
    return blobs

def plot_detected_blobs(mask, blobs, title):
    fig, ax = plt.subplots(figsize=(6,6))
    ax.imshow(mask, cmap='gray')
    ax.set_title(title)
    ax.axis("off")

    for blob in blobs:
        y, x, r = blob
        circle =plt.Circle((x,y), r, color = 'red', linewidth = 2, fill = False)
        ax.add_patch(circle)

    plt.show()
class Plotter:

    @staticmethod
    def plot_image_channels(image, red, green, blue):
        fig, axes = plt.subplots(2, 2, figsize=(12,8))

        axes[0,0].imshow(image)
        axes[0,0].set_title("Original picture")
        axes[0,0].axis("off")

        axes[0,1].imshow(red,cmap = "Reds")
        axes[0,1].set_title("Red channel")
        axes[0,1].axis("off")

        axes[1,0].imshow(green, cmap = "Greens")
        axes[1,0].set_title("Green channel")
        axes[1,0].axis("off")    

        axes[1,1].imshow(blue, cmap = "Blues")
        axes[1,1].set_title("blue channel")
        axes[1,1].axis("off")    

        plt.show()
    @staticmethod
    def plot_histograms(red, green ,blue):
        fig, axes = plt.subplots(1,3, figsize =(15,5))

        axes[0].hist(red.ravel(), bins=256, color="red", alpha= 0.7)
        axes[0].set_title('Red Channel Histogram')
        axes[0].set_xlabel('Pixel Intensity')
        axes[0].set_ylabel('Frequency')

        axes[1].hist(green.ravel(), bins=256, color="green", alpha= 0.7)
        axes[1].set_title('Green Channel Histogram')
        axes[1].set_xlabel('Pixel Intensity')
        axes[1].set_ylabel('Frequency')

        axes[2].hist(blue.ravel(), bins=256, color="blue", alpha= 0.7)
        axes[2].set_title('Blue Channel Histogram')
        axes[2].set_xlabel('Pixel Intensity')
        axes[2].set_ylabel('Frequency')

        plt.show()

    @staticmethod
    def plot_filtered_image(original, filtered):
        fig, axes = plt.subplots(1, 2, figsize=(12,6))

        axes[0].imshow(original)
        axes[0].set_title("Original Image")
        axes[0].axis("off")

        axes[1].imshow(filtered)
        axes[1].set_title("Filtered Image (Median Blur)")
        axes[1].axis("off")

        plt.show()

    @staticmethod
    def plot_color_masks(image, red_mask, yellow_mask, green_mask):
        fig, axes = plt.subplots(1, 3, figsize = (18,6))

        axes[0].imshow(red_mask, cmap='gray')
        axes[0].set_title("Red Mask")
        axes[0].axis("off")

        axes[1].imshow(yellow_mask, cmap='gray')
        axes[1].set_title("Yellow Mask")
        axes[1].axis("off")

        axes[2].imshow(green_mask, cmap='gray')
        axes[2].set_title("Green Mask")
        axes[2].axis("off")

        plt.show()
    
    @staticmethod
    def plot_blobs(red_mask, yellow_mask, green_mask, red_blobs, yellow_blobs, green_blobs):
        fig, axes = plt.subplots(1, 3, figsize=(18,6))

        axes[0].imshow(red_mask, cmap ='gray')
        axes[0].set_title("Red blobs")
        axes[0].axis("off")
        for blob in red_blobs:
            y, x, r = blob
            circle = plt.Circle((x,y), r, color ='red', linewidth = 2, fill = False)
            axes[0].add_patch(circle)

        axes[1].imshow(yellow_mask, cmap ='gray')
        axes[1].set_title("Yellow blobs")
        axes[1].axis("off")
        for blob in yellow_blobs:
            y, x, r = blob
            circle = plt.Circle((x,y), r, color ='yellow', linewidth = 2, fill = False)
            axes[1].add_patch(circle)

        axes[2].imshow(green_mask, cmap ='gray')
        axes[2].set_title("Green blobs")
        axes[2].axis("off")
        for blob in green_blobs:
            y, x, r = blob
            circle = plt.Circle((x,y), r, color ='green', linewidth = 2, fill = False)
            axes[2].add_patch(circle)

        plt.show()

image_path = "traffic_light_image_1.png"
processor = ImageProcessor(image_path)
red, green, blue = processor.get_color_channels()

red_mask, yellow_mask, green_mask = processor.get_color_masks()

red_blobs = detect_blobs(red_mask)
yellow_blobs = detect_blobs(yellow_mask)
green_blobs = detect_blobs(green_mask)

print(f"Red Blobs: {len(red_blobs)}")
print(f"Yellow Blobs: {len(yellow_blobs)}")
print(f"Green Blobs: {len(green_blobs)}")

# snr_red = processor.calculate_snr(red)
# snr_green = processor.calculate_snr(green)
# snr_blue = processor.calculate_snr(blue)

#print(f"SNR for Red Channel {snr_red:.2f}")
#print(f"SNR for Green Channel {snr_green:.2f}")
#print(f"SNR for Blue Channel {snr_blue:.2f}")

# filtered_image = processor.apply_median_filter(kernel_size = 5)


#Plotter.plot_image_channels(processor.image_rgb,red,green,blue)
#Plotter.plot_histograms(red, green, blue)
#Plotter.plot_filtered_image(processor.image_rgb, filtered_image)
#Plotter.plot_color_masks(processor.image_rgb, red_mask, yellow_mask, green_mask)
Plotter.plot_blobs(red_mask, yellow_mask, green_mask, red_blobs, yellow_blobs, green_blobs)