|
|
from Data.config import * |
|
|
import logging |
|
|
|
|
|
|
|
|
class SortModel: |
|
|
def __init__(self, model, classes_to_delete=classes_to_delete, conf_threshold=det_conf, expected_segments=expected_segments): |
|
|
self.model = model |
|
|
self.classes_to_delete = classes_to_delete |
|
|
self.conf_threshold = conf_threshold |
|
|
self.expected_segments = expected_segments |
|
|
|
|
|
def get_center(self, detection): |
|
|
_, x_min, y_min, x_max, y_max = detection |
|
|
center_x = (x_min + x_max) / 2 |
|
|
center_y = (y_min + y_max) / 2 |
|
|
return center_x, center_y |
|
|
|
|
|
def sort_and_group_detections(self, detections): |
|
|
detections_with_centers = [(d[0], *self.get_center(d)) for d in detections] |
|
|
sorted_detections = sorted(detections_with_centers, key=lambda x: (x[1], x[2])) |
|
|
|
|
|
if not sorted_detections: |
|
|
return [] |
|
|
|
|
|
threshold_x = (sorted_detections[0][1] * 0.5) |
|
|
rows = [] |
|
|
current_row = [] |
|
|
current_x = sorted_detections[0][1] |
|
|
|
|
|
for detection in sorted_detections: |
|
|
class_name, center_x, center_y = detection |
|
|
|
|
|
if abs(center_x - current_x) > threshold_x: |
|
|
rows.append(sorted(current_row, key=lambda x: x[2])) |
|
|
current_row = [] |
|
|
current_x = center_x |
|
|
|
|
|
current_row.append(detection) |
|
|
|
|
|
if current_row: |
|
|
rows.append(sorted(current_row, key=lambda x: x[2])) |
|
|
|
|
|
max_columns = max(len(row) for row in rows) if rows else 0 |
|
|
grid_matrix = [] |
|
|
for row in rows: |
|
|
grid_row = [d[0] for d in row] |
|
|
grid_row.extend([''] * (max_columns - len(row))) |
|
|
grid_matrix.append(grid_row) |
|
|
|
|
|
transposed_matrix = list(map(list, zip(*grid_matrix))) |
|
|
return transposed_matrix |
|
|
|
|
|
@staticmethod |
|
|
def sequence(matrix, expected_segments): |
|
|
if not expected_segments: |
|
|
return False |
|
|
|
|
|
for sequence in matrix: |
|
|
segment_index = 0 |
|
|
for item in sequence: |
|
|
if item == expected_segments[segment_index]: |
|
|
continue |
|
|
elif segment_index < len(expected_segments) - 1 and item == expected_segments[segment_index + 1]: |
|
|
segment_index += 1 |
|
|
else: |
|
|
return False |
|
|
if segment_index != len(expected_segments) - 1: |
|
|
return False |
|
|
return True |
|
|
|
|
|
def process_image(self, image_path, predicted_class): |
|
|
|
|
|
results = self.model(image_path) |
|
|
|
|
|
planogram_ghw_count = 0 |
|
|
planogram_blanks_count = 0 |
|
|
planogram_ghw = '' |
|
|
planogram_valid_sequence = '' |
|
|
yolo_detections = [] |
|
|
|
|
|
for result in results: |
|
|
for box in result.boxes: |
|
|
class_name = self.model.names[int(box.cls[0])] |
|
|
|
|
|
if box.conf[0] >= self.conf_threshold: |
|
|
if class_name in self.classes_to_delete: |
|
|
planogram_ghw_count += 1 |
|
|
continue |
|
|
|
|
|
if box.conf[0] >= self.conf_threshold: |
|
|
x_min, y_min, x_max, y_max = box.xyxy[0] |
|
|
yolo_detections.append((class_name, x_min.item(), y_min.item(), x_max.item(), y_max.item())) |
|
|
|
|
|
planogram_blanks_count = len(yolo_detections) |
|
|
|
|
|
grid_matrix = self.sort_and_group_detections(yolo_detections) |
|
|
|
|
|
|
|
|
print("Grid Matrix:") |
|
|
for row in grid_matrix: |
|
|
print(row) |
|
|
print("\n") |
|
|
|
|
|
if planogram_blanks_count == planogram_ghw_count: |
|
|
planogram_ghw = "yes" |
|
|
else: |
|
|
planogram_ghw = "no" |
|
|
|
|
|
planogram_valid_sequence = self.sequence(grid_matrix, self.expected_segments) |
|
|
if planogram_valid_sequence: |
|
|
planogram_valid_sequence = "yes" |
|
|
else: |
|
|
planogram_valid_sequence = "no" |
|
|
|
|
|
return { |
|
|
"Planogram Blanks Count": planogram_blanks_count, |
|
|
"Planogram GHW Count": planogram_ghw_count, |
|
|
"Planogram Valid Sequence": planogram_valid_sequence, |
|
|
"Planogram GHW": planogram_ghw, |
|
|
"Planogram Name": predicted_class, |
|
|
|
|
|
"POSM Blanks Count": "N/A", |
|
|
"POSM GHW Count": "N/A", |
|
|
"POSM Valid Sequence": "N/A", |
|
|
"POSM GHW": "N/A", |
|
|
"POSM Name": "N/A", |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
