Utils/segment.py

import cv2
import os
import uuid
import numpy as np
from Data.config import *

class CropModel:

    def __init__(self, model, model_input_shape=(frame_shape[0], frame_shape[1], 3)):
        self.model = model
        self.model_input_shape = model_input_shape
        self.conf_threshold = seg_conf

    def image_prediction_mask(self, image):
        predict = self.model.predict(image)

        # Extract the masks and move to CPU only if necessary
        mask_tensor = predict[0].masks.data[0]  # Assuming this is the mask tensor
        if mask_tensor.is_cuda:
            mask = mask_tensor.cpu().numpy() * 255  # Convert to NumPy array after moving to CPU
        else:
            mask = mask_tensor.numpy() * 255  # No need to move to CPU if it's already on CPU

        mask = mask.astype("uint8")

        # Extract class IDs and class names from the predictions
        class_ids_tensor = predict[0].boxes.cls
        if class_ids_tensor.is_cuda:
            class_ids = class_ids_tensor.cpu().numpy()  # Move to CPU
        else:
            class_ids = class_ids_tensor.numpy()

        class_names = [self.model.names[int(cls_id)] for cls_id in class_ids]  # Convert indices to names

        # If any class is found, return the first one
        if class_names:
            class_name = class_names[0]  # Return the first detected class
        else:
            class_name = "Not Found"  # If no class is detected

        return mask, class_name

    
    @staticmethod
    def get_mask_corner_points(mask):
        _, thresh = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

        if not contours:
            return None

        cnt = max(contours, key=cv2.contourArea)
        cnt_approx = cv2.approxPolyDP(cnt, 0.03 * cv2.arcLength(cnt, True), True)

        return cnt_approx.reshape((4, 2)) if len(cnt_approx) == 4 else None

    @staticmethod
    def get_order_points(points):
        rect = np.zeros((4, 2), dtype="float32")
        s = points.sum(axis=1)
        diff = np.diff(points, axis=1)

        rect[0] = points[np.argmin(s)]
        rect[2] = points[np.argmax(s)]
        rect[1] = points[np.argmin(diff)]
        rect[3] = points[np.argmax(diff)]

        return rect

    @staticmethod
    def expand_bounding_box(points, expand_ratio):
        center = np.mean(points, axis=0)
        expanded_points = points + (points - center) * expand_ratio
        return expanded_points.astype("float32")

    def point_transform(self, image, points):
        ordered_points = self.get_order_points(points)
        expanded_points = self.expand_bounding_box(ordered_points, EXPAND_RATIO)
        height, width = image.shape[:2]

        dst = np.array([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]], dtype="float32")
        M = cv2.getPerspectiveTransform(expanded_points, dst)
        warped_image = cv2.warpPerspective(image, M, (width, height))

        return warped_image

    @staticmethod
    def add_padding_to_image(image, padding_size):
        return cv2.copyMakeBorder(image, padding_size, padding_size, padding_size, padding_size, cv2.BORDER_CONSTANT, value=[0, 0, 0])

    # def get_predicted_warped_image(self, image):
    #     mask, class_name = self.image_prediction_mask(image)
    #     corner_points = self.get_mask_corner_points(mask)

    #     if corner_points is None:
    #         return None, class_name

    #     warped_image = self.point_transform(image, corner_points)
    #     padded_image = self.add_padding_to_image(warped_image, OUTER_PADDING_SIZE)


    #     return padded_image, class_name

    def get_predicted_warped_image(self, image, save_dir="store"):
        mask, class_name = self.image_prediction_mask(image)
        corner_points = self.get_mask_corner_points(mask)

        if corner_points is None:
            return None, class_name

        # Perform the transformation
        warped_image = self.point_transform(image, corner_points)

        # Add padding to the image
        padded_image = self.add_padding_to_image(warped_image, OUTER_PADDING_SIZE)

        # Ensure the directory exists
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        # Generate a random image name
        random_name = f"image_{uuid.uuid4().hex[:8]}.jpg"  # First 8 characters of a UUID
        save_path = os.path.join(save_dir, random_name)

        # Save the padded image
        cv2.imwrite(save_path, padded_image)
        print(f"Image saved at: {save_path}")

        return padded_image, class_name