File size: 4,637 Bytes
cb90fd0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
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
|