Upload New File

Process saliency map/ObtainGrid.py

+import cv2
+import math
+import numpy as np
+from shapely.geometry import Polygon
+
+# Constants
+CELL_SIZE = 128
+CELL_AREA = CELL_SIZE * CELL_SIZE
+SEED_THRESHOLD = 70
+LEVEL2_MULTIPLIER = 2
+DIRECTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)] 
+
+def convert_to_color(grayscale_image):
+    """Convert a grayscale image to RGB."""
+    return cv2.cvtColor(grayscale_image, cv2.COLOR_GRAY2RGB)
+
+class QpGrid:
+    def __init__(self, path, qp_grid_path, qp_target, img_size, frame_count, qp_max_difference=8, verbose=0):
+        print('Initializing QpGrid...')
+        self.path = path
+        self.qp_grid_path = qp_grid_path
+        self.verbose = verbose
+        self.qp_target = qp_target
+        self.img_size = img_size
+        self.qp_max_difference = qp_max_difference
+        self.levels = self.level_to_qp()
+        self.qp_lines = [""] * frame_count
+        self.idx = 0
+
+        width, height = img_size
+        self.grid_width = math.ceil(width / CELL_SIZE)
+        self.grid_height = math.ceil(height / CELL_SIZE)
+
+        print("Level 2 multiplier:", LEVEL2_MULTIPLIER)
+
+    def level_to_qp(self):
+        """Map levels to QP values."""
+        return np.array([self.qp_max_difference, int(self.qp_max_difference / 2), 0])
+
+    def find_attention_points(self, image):
+        """Automatically find seed points in the image."""
+        _, binary_image = cv2.threshold(image, SEED_THRESHOLD, 255, cv2.THRESH_BINARY)
+        contours, _ = cv2.findContours(binary_image, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+
+        centroids = []
+        if len(contours) == 0:
+            return centroids, [], []
+
+        if len(contours) == 1:
+            return centroids, [], contours
+
+        combined_contour = np.concatenate(contours)
+        hull = cv2.convexHull(combined_contour)
+
+        return centroids, hull, contours
+
+    def grid_to_qp(self, grid):
+        """Convert grid levels to QP values."""
+        return self.qp_target + self.levels[grid]
+
+    def save_qp_line(self, grid, idx):
+        """Save a single line of QP values."""
+        self.qp_lines[idx] = ', '.join(str(grid[i, j]) for j in range(grid.shape[1]) for i in range(grid.shape[0])) + ","
+
+    def save_qp_to_file(self):
+        """Save all QP lines to a file."""
+        with open(self.path, 'w') as f:
+            f.write('\n'.join(self.qp_lines))
+
+    def display_attention_points(self, image, seeds):
+        """Display seed points on the image."""
+        for idx, seed in enumerate(seeds):
+            cv2.circle(image, seed, 5, (0, 0, 255), -1)
+            cv2.putText(image, f'Seed {idx+1}', (seed[0]+10, seed[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
+
+    def display_grid(self, image, grid):
+        """Overlay the grid and its values on the image."""
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        font_scale = 1
+        font_color = (255, 255, 255)
+        line_type = 2
+
+        colors = {1: (128, 0, 128), 2: (0, 0, 255), 3: (0, 165, 255)}
+
+        for i in range(self.grid_height):
+            for j in range(self.grid_width):
+                grid_value = 3 - grid[i, j]
+                color = colors.get(grid_value, (255, 255, 255))
+                cv2.rectangle(image, (j * CELL_SIZE, i * CELL_SIZE), ((j + 1) * CELL_SIZE, (i + 1) * CELL_SIZE), color, 2)
+                cv2.putText(image, str(grid_value), (j * CELL_SIZE + 10, i * CELL_SIZE + 30), font, font_scale, font_color, line_type)
+
+    def assign_seeds_to_level1(self, seeds, grid):
+        """Assign level 1 to cells containing seeds."""
+        for seed in seeds:
+            cell = (seed[0] // CELL_SIZE, seed[1] // CELL_SIZE)
+            try:
+                grid[cell[1], cell[0]] = 2
+            except IndexError:
+                print("Error with seed:", cell)
+        return grid
+
+    def assign_level2_to_grid(self, grid, level1_cells):
+        """Assign level 2 to neighboring cells of level 1 cells."""
+        update_list = list(level1_cells)
+        level2_cell_count = len(level1_cells)
+        level1_cell_count = 0
+
+        while level1_cell_count < level2_cell_count * LEVEL2_MULTIPLIER:
+            new_update_list = []
+            for cell in update_list:
+                i, j = cell
+                for dx, dy in DIRECTIONS:
+                    ni, nj = i + dx, j + dy
+                    if 0 <= ni < self.grid_height and 0 <= nj < self.grid_width and grid[ni, nj] == 0:
+                        grid[ni, nj] = 1
+                        new_update_list.append((ni, nj))
+            level1_cell_count += len(update_list)
+            update_list = new_update_list
+            if self.verbose >= 3:
+                print("Level 1 count:", level1_cell_count, "Level 2 target:", level2_cell_count * LEVEL2_MULTIPLIER)
+            if not update_list:
+                if self.verbose >= 3:
+                    print("No more cells to update.")
+                break
+
+        return grid
+
+    def assign_contours_to_level1(self, contours, hull, grid):
+        """Assign level 1 to cells intersecting with contours or hull."""
+        cell_polygons = {
+            (i, j): Polygon([(j * CELL_SIZE, i * CELL_SIZE), ((j + 1) * CELL_SIZE, i * CELL_SIZE),
+                             ((j + 1) * CELL_SIZE, (i + 1) * CELL_SIZE), (j * CELL_SIZE, (i + 1) * CELL_SIZE)])
+            for i in range(self.grid_height) for j in range(self.grid_width)
+        }
+
+        level1_cells = set()
+
+        for contour in contours:
+            try:
+                if len(contour) < 4:
+                    if self.verbose >= 1:
+                        print("Skipping contour with fewer than 4 points:", contour)
+                    continue
+                contour_polygon = Polygon(contour[:, 0, :])
+                for (i, j), cell_polygon in cell_polygons.items():
+                    if (i, j) in level1_cells:
+                        continue
+                    intersection = cell_polygon.intersection(contour_polygon)
+                    if intersection.area / CELL_AREA >= 0.01:
+                        grid[i, j] = 2
+                        level1_cells.add((i, j))
+            except Exception as e:
+                if self.verbose >= 1:
+                    print("Error processing contour:", e)
+
+        if hull:
+            hull_polygon = Polygon(hull[:, 0, :])
+            for (i, j), cell_polygon in cell_polygons.items():
+                if (i, j) in level1_cells:
+                    continue
+                intersection = cell_polygon.intersection(hull_polygon)
+                if intersection.area / CELL_AREA >= 0.25:
+                    grid[i, j] = 2
+                    level1_cells.add((i, j))
+
+        return grid, level1_cells
+
+    def process_image(self, image, idx):
+        """Process a single image to generate the QP grid."""
+        if self.verbose >= 3:
+            print(f"Processing image {idx}...")
+
+        grid = np.zeros((self.grid_height, self.grid_width), dtype=int)
+
+        seeds, hull, contours = self.find_attention_points(image)
+        grid, level1_cells = self.assign_contours_to_level1(contours, hull, grid)
+        grid = self.assign_level2_to_grid(grid, level1_cells)
+
+        qp = self.grid_to_qp(grid)
+        self.save_qp_line(qp, idx)
+
+        if self.verbose >= 2:
+            image = convert_to_color(image)
+            cv2.drawContours(image, contours, -1, (255, 255, 0), 2)
+            if len(contours) > 1 and hull is not None:
+                try:
+                    cv2.drawContours(image, [hull], -1, (255, 0, 0), 2)
+                except Exception as e:
+                    print(e)
+
+            self.display_grid(image, grid.astype(int))
+            image = cv2.resize(image, (1920, 1080))
+            cv2.imwrite(f'{self.qp_grid_path}/image_{idx}.png', image)
\ No newline at end of file