How good are you at Vacuuming? Let's find out.
import cv2
import numpy as np
import yaml
from roboflow import Roboflow
import datetime
from vcTracking import HeatmapGenerator
from vcVideoHandler import VideoHandler, FrameOverlay

# Read config file
def read_config(file_path):
    with open(file_path, 'r') as file:
        config_dict = yaml.safe_load(file)
    return config_dict

def concatenate_images(img_arrays, titles, m, n):
    """
    Concatenate images into a grid of m rows and n columns with respective titles.
    
    Parameters:
    - img_arrays: List of image arrays.
    - titles: List of image titles.
    - m: Number of rows.
    - n: Number of columns.

    Returns:
    - concatenated image
    """
    
    if len(img_arrays) != m*n:
        raise ValueError("The number of images provided does not match m x n.")
    
    reference_height, reference_width = img_arrays[0].shape[:2]
    
    # Resize images based on the first image's dimensions
    images = [cv2.resize(img, (reference_width, reference_height)) for img in img_arrays]
    
    # Draw titles on the images
    font_scale = 1.5  # Increased font size to 3 times
    font_thickness = 2
    for idx, title in enumerate(titles):
        (text_width, text_height), _ = cv2.getTextSize(title, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
        
        # Calculating the starting coordinates for centered text
        start_x = (reference_width - text_width) // 2
        start_y = text_height + 10

        cv2.rectangle(images[idx], (0, 0), (reference_width, text_height + 20), (255, 255, 255), -1)
        cv2.putText(images[idx], title, (start_x, start_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), font_thickness)
    
    # Concatenate images in m x n format
    rows = []
    for i in range(0, len(images), n):
        rows.append(np.hstack(images[i:i+n]))
    
    concatenated_img = np.vstack(rows)
    
    return concatenated_img

# Setup Robboflow Stuff
def setup_roboflow(config):
    # Set up Roboflow Model (using Hosted API now)
    rf = Roboflow(api_key=config['roboflow']['api_key'])
    project = rf.workspace().project(config['roboflow']['model'])
    model = project.version(config['roboflow']['version']).model

    # Retrun model
    return model

# Run roboflow inference
def run_inference(model, frame):
    # Run inference on the frame
    results = model.predict(frame,
                    confidence=60, 
                    overlap=30).json()
    
    return results

# Main
if __name__ == "__main__":
    # Read config file
    config_path = 'config.yaml'  # you can modify this to the path of your config file
    config = read_config(config_path)
    # Show content of config file
    print(config)  

    # Init Videohandler with video path
    video = VideoHandler(config['video']['filename'])

    # Get first frame
    frame = video.get_first_frame()
    overlay = FrameOverlay(frame)

    # Have user draw polygon
    message = 'Draw a polygon border where you will be vacuuming: When done press "d"'
    overlay.center_text_at_top(message)
    points = overlay.draw_polygon(window_name='First Frame')

    # Show user the polygon
    overlay.set_frame(frame)
    overlay.draw_polygon(points)
    prompt_text = "Is this what you want? (y/n)"
    overlay.center_text_at_top(prompt_text, background=True)
    cv2.imshow('First Frame', overlay.get_frame())

    # Have user confirm polygon
    while True:
        key = cv2.waitKey(1) & 0xFF
        if key == ord("y"):
            break
        elif key == ord("n"):
            overlay.set_frame(frame)
            points = overlay.draw_polygon('First Frame')
            overlay.set_frame(frame)
            overlay.draw_polygon(points)
            overlay.center_text_at_top(prompt_text, background=True)
            cv2.imshow('First Frame', overlay.get_frame())

    # Destroy window
    cv2.destroyAllWindows()

    # Configure Roboflow model
    model = setup_roboflow(config)

    # Init heatmap / path
    heatmap_gen = HeatmapGenerator(generate_heatmap=True, generate_path=True, generate_binary_mask=True, polygon=points)

    # Loop through video
    heatmap_frame = None
    path_frame = None
    binary_mask = None
    percentage_complete = 0
    out = None
    while True:
        # Get next frame
        frame = video.get_next_frame()
        # Check if frame is valid
        if frame is None:
            break

        # Inference on frame
        results_cv = run_inference(model, frame)
        # Extract results
        results = []
        for prediction in results_cv['predictions']:
            # Results
            results.append([prediction['x']-prediction['width']/2, prediction['y']-prediction['height']/2, prediction['x']+prediction['width']/2, prediction['y']+prediction['height']/2, prediction['class']])

        # Overlay polygon on frame
        overlay = FrameOverlay(frame)
        overlay.draw_polygon(points)

        # Display heatmap / path
        for result in results:
            # Check if result is vacuum:
            if result[-1] != config['video']['object_track']:
                continue

            # Get bounding box
            bbox = result[:4]  # assuming you have bounding box coordinates like this
            # Overlay bbox and label
            overlay.draw_bbox_with_label(bbox, result[-1])

            # Update
            heatmap_frame, path_frame, binary_mask = heatmap_gen.process_frame(frame, bbox)
            # Calculate untouched percentage (this needs to be on only if heatmap is on)
            percentage_complete = 100.0 - heatmap_gen.calculate_untouched_percentage()
            # print(f"Percentage of untouched pixels inside the polygon: {percentage:.2f}%")
            # cv2.imshow('heatmap_frame', heatmap_frame)
            # cv2.imshow('path_frame', path_frame)
            # cv2.imshow('mask_frame', binary_mask)

        # Show frame
        overlay_frame = overlay.get_frame()
        # cv2.imshow('Video', overlay_frame)

        # Concatenate images
        if heatmap_frame is None or path_frame is None or binary_mask is None:
            heatmap_frame = np.zeros_like(frame)
            path_frame = np.zeros_like(frame)
            binary_mask = np.zeros_like(frame)
        concatenated = concatenate_images([overlay_frame, heatmap_frame, path_frame, binary_mask], ["Frame", "Heatmap", "Path", "Coverage: {:.2f}%".format(percentage_complete)], 1, 4)
        cv2.imshow("Concatenated Images", concatenated)

        # Write to video
        if config["output"]["save_output"]:
            if out is None:
                fps = 15
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # or use 'XVID' 

                # Get timestamp to append to filenae    
                current_timestamp = datetime.datetime.now()
                formatted_timestamp = current_timestamp.strftime('%Y%m%d_%H%M%S')
                filename_save = config["output"]["save_path"].split('.')[0] + '_' + formatted_timestamp + '.mp4'
                # Create video writer
                out = cv2.VideoWriter(filename_save, fourcc, fps, (concatenated.shape[1], concatenated.shape[0]))
            out.write(concatenated)

        # Press 'q' to quit
        if cv2.waitKey(30) & 0xFF == ord('q'):
            break

    # Release video and close
    out.release()
    video.release()
    cv2.destroyAllWindows()
scripts/main.py
reedajohns/vacuum-challenge
import cv2
import numpy as np
import yaml
from roboflow import Roboflow
import datetime

# Heatmap / Path Class
class HeatmapGenerator:
    def __init__(self, generate_heatmap=True, generate_path=True, generate_binary_mask=True, polygon=None):
        self.heatmap_width = 0
        self.heatmap_height = 0
        self.accumulator = None
        self.path = None
        self.object_centroids = []
        self.binary_mask = None
        # Which outputs to generate
        self.generate_heatmap = generate_heatmap
        self.generate_path = generate_path
        self.generate_binary_mask = generate_binary_mask
        # Polygon
        self.polygon = None if polygon is None else np.array(polygon, dtype=np.int32)  # Convert the polygon points to numpy array or default to None

    @staticmethod
    def overlay_image_alpha(img, heatmap, alpha=0.5):
        img_float = img.astype(np.float32) / 255.0
        heatmap_float = heatmap.astype(np.float32) / 255.0
        blended = cv2.addWeighted(img_float, 1 - alpha, heatmap_float, alpha, 0)
        return (blended * 255).astype(np.uint8)

    def init_heatmap_path(self, frame):
        self.heatmap_width = frame.shape[1]
        self.heatmap_height = frame.shape[0]
        self.accumulator = np.zeros((self.heatmap_height, self.heatmap_width), dtype=np.float32)
        self.path = np.zeros((self.heatmap_height, self.heatmap_width, 3), dtype=np.uint8)
        # Initialize binary_mask if the feature is turned on (Do this to others)
        if self.generate_binary_mask:
            self.binary_mask = np.zeros((self.heatmap_height, self.heatmap_width, 3), dtype=np.uint8)
            self.binary_mask[:] = (128, 128, 128)  # Start with gray outside

            # If no polygon was provided, consider the entire frame as the polygon
            if self.polygon is None:
                self.polygon = np.array([
                    [0, 0],
                    [0, self.heatmap_height-1],
                    [self.heatmap_width-1, self.heatmap_height-1],
                    [self.heatmap_width-1, 0]
                ], dtype=np.int32)

            # Fill polygon area with green  
            cv2.fillPoly(self.binary_mask, [self.polygon], (0, 255, 0))

    def update_accumulator(self, x0, y0, x1, y1):
        centroid = ((x0 + x1) // 2, (y0 + y1) // 2)
        self.object_centroids.append(centroid)
    
        # Ensure bounding box values are integers
        self.accumulator[y0:y1, x0:x1] += 1

        # Update binary mask
        if self.generate_binary_mask:
            # Set the pixels inside the bounding box and inside the polygon to black
            mask_inside_polygon = np.zeros_like(self.binary_mask[:, :, 0], dtype=np.uint8)
            cv2.fillPoly(mask_inside_polygon, [self.polygon], 255)
            mask = (self.binary_mask[y0:y1, x0:x1] == [0, 255, 0]).all(axis=2) & (mask_inside_polygon[y0:y1, x0:x1] == 255)
            self.binary_mask[y0:y1, x0:x1][mask] = [0, 0, 0]

            # Set the pixels inside the bounding box and outside the polygon to red
            mask_outside_polygon = (self.binary_mask[y0:y1, x0:x1] == [128, 128, 128]).all(axis=2)
            self.binary_mask[y0:y1, x0:x1][mask_outside_polygon] = [0, 0, 255]

            # cv2.imshow('mask', self.binary_mask)
            # cv2.waitKey(1)
            # input('bru')

    def draw_path(self, frame):
        if not self.generate_path:
            return None

        path_frame = frame.copy()
        num_centroids = len(self.object_centroids)
        for i in range(1, num_centroids):
            alpha = i / num_centroids
            color = (int(255 * (1-alpha)), 0, int(255 * alpha))
            cv2.line(path_frame, self.object_centroids[i-1], self.object_centroids[i], color, 6)
        return path_frame
    
    def calculate_untouched_percentage(self):
        # 1. Create a binary mask for the polygon area
        polygon_mask = np.zeros_like(self.binary_mask[:, :, 0], dtype=np.uint8)
        cv2.fillPoly(polygon_mask, [self.polygon], 1)
        
        # 2. Extract the green channel of our binary_mask
        green_channel = self.binary_mask[:, :, 1]
        
        # 3. Multiply the polygon_mask with the green channel
        green_inside_polygon = polygon_mask * green_channel
        
        # Count untouched (green) pixels and total pixels inside the polygon
        untouched_pixel_count = np.sum(green_inside_polygon == 255)
        total_pixel_count = np.sum(polygon_mask == 1)
        
        # 4. Calculate the percentage
        if total_pixel_count == 0:
            return 0  # Avoid division by zero
        untouched_percentage = (untouched_pixel_count / total_pixel_count) * 100
        
        return untouched_percentage

    def process_frame(self, frame, bbox):
        if self.accumulator is None:
            self.init_heatmap_path(frame)

        x0, y0, x1, y1 = bbox
        # Ensure bounding box values are integers
        x0, y0, x1, y1 = map(int, [x0, y0, x1, y1])
        self.update_accumulator(x0, y0, x1, y1)

        heatmap_frame = None
        if self.generate_heatmap:
            heatmap = cv2.applyColorMap(np.uint8(255 * self.accumulator / self.accumulator.max()), cv2.COLORMAP_JET)
            heatmap_frame = self.overlay_image_alpha(frame, heatmap, alpha=0.7)

        path_frame = self.draw_path(frame)

        # Return binary mask frame or None depending on the toggle
        binary_mask_frame = self.binary_mask if self.generate_binary_mask else None

        return heatmap_frame, path_frame, binary_mask_frame
scripts/vcTracking.py
reedajohns/vacuum-challenge
import cv2
import numpy as np
import yaml
from roboflow import Roboflow
import datetime

class VideoHandler:
    def __init__(self, video_path):
        self.cap = cv2.VideoCapture(video_path)
        if not self.cap.isOpened():
            raise ValueError("Unable to open video file.")
        self.current_frame = None

    def get_first_frame(self):
        success, frame = self.cap.read()
        if success:
            self.current_frame = frame
            return frame
        else:
            return None

    def get_next_frame(self):
        success, frame = self.cap.read()
        if success:
            self.current_frame = frame
            return frame
        return None

    def release(self):
        self.cap.release()

class FrameOverlay:
    def __init__(self, frame):
        self.frame = frame.copy()
        self.frame_original = frame.copy()

    def draw_polygon(self, points=None, window_name='Draw Polygon'):
        if points is None:
            return self._interactive_draw_polygon(window_name)
        else:
            if len(points) > 1:
                cv2.polylines(self.frame, [np.array(points)], isClosed=True, color=(255, 50, 0), thickness=6)
                self._shade_polygon(points)

    def _interactive_draw_polygon(self, window_name):
    # def draw_polygon(self, window_name='Draw Polygon'):
        def mouse_event(event, x, y, flags, param):
            nonlocal drawing
            if event == cv2.EVENT_LBUTTONDOWN:
                drawing = "True"
                points.append((x, y))
                if len(points) > 1:
                    cv2.line(self.frame, points[-2], points[-1], (255, 50, 0), 6)
                cv2.circle(self.frame, (x, y), 8, (255, 0, 0), -1)
                cv2.imshow(window_name, self.frame)
        
        points = []
        drawing = False
        cv2.namedWindow(window_name)
        cv2.setMouseCallback(window_name, mouse_event)
        
        while True:
            cv2.imshow(window_name, self.frame)
            key = cv2.waitKey(1) & 0xFF
            if key == ord("d"):
                if len(points) > 1:
                    cv2.line(self.frame, points[-1], points[0], (0, 0, 255), 2)
                break
        cv2.destroyAllWindows()
        return points
    
    def _shade_polygon(self, points, shading_color=(128, 128, 128), opacity=0.4):
        # Create a blank image
        overlay = np.zeros_like(self.frame)
        
        # Draw the filled polygon on the blank image
        cv2.fillPoly(overlay, [np.array(points)], shading_color)
        
        # Blend the original frame and the shaded polygon image
        cv2.addWeighted(overlay, opacity, self.frame, 1 - opacity, 0, self.frame)

    def center_text_at_top(self, text, font=cv2.FONT_HERSHEY_SIMPLEX, scale=1, color=(255, 0, 0), thickness=2, background=True):
        (text_width, text_height), _ = cv2.getTextSize(text, font, scale, thickness)
        while text_width > self.frame.shape[1]:
            scale -= 0.1
            (text_width, text_height), _ = cv2.getTextSize(text, font, scale, thickness)
    
        start_x = (self.frame.shape[1] - text_width) // 2
        start_y = text_height + 10
        
        if background:
            cv2.rectangle(self.frame, (start_x - 5, start_y - text_height - 5), (start_x + text_width + 5, start_y + 5), (255, 255, 255), -1)
        
        cv2.putText(self.frame, text, (start_x, start_y), font, scale, color, thickness)

    def draw_bbox_with_label(self, bbox, label, box_color=(0, 255, 0), text_color=(255, 255, 255), font=cv2.FONT_HERSHEY_SIMPLEX, font_scale=0.5, font_thickness=1):
        """
        Draw a bounding box and label on the image frame.

        Parameters:
        - bbox: Tuple of (x0, y0, x1, y1) where (x0, y0) is the top-left coordinate and (x1, y1) is the bottom-right coordinate of the box.
        - label: String label to be drawn above the bounding box.
        - box_color: Color of the bounding box. Default is green.
        - text_color: Color of the label text. Default is white.
        - font: Font to be used for the label text. Default is cv2.FONT_HERSHEY_SIMPLEX.
        - font_scale: Font scale factor. Default is 0.5.
        - font_thickness: Thickness of the font. Default is 1.
        """
        
        # Draw bounding box
        x0, y0, x1, y1 = bbox
        x0, y0, x1, y1 = map(int, [x0, y0, x1, y1])
        cv2.rectangle(self.frame, (x0, y0), (x1, y1), box_color, 2)

        # Calculate the size of the text
        (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, font_thickness)

        # Draw a filled rectangle behind the text for better visibility
        cv2.rectangle(self.frame, (x0, y0 - text_height - 5), (x0 + text_width + 5, y0), box_color, -1)

        # Draw label above the bounding box
        cv2.putText(self.frame, label, (x0, y0 - 5), font, font_scale, text_color, font_thickness)

    def get_frame(self):
        return self.frame

    def set_frame(self, frame):
        self.frame = frame.copy()
scripts/vcVideoHandler.py
reedajohns/vacuum-challenge