Spaces:
Running
Running
| import json | |
| import numpy as np | |
| import cv2 | |
| from shapely.geometry import Polygon | |
| from io import BytesIO | |
| from ultralytics import YOLO | |
| import torch | |
| # Define colors | |
| COLORS = [(98, 231, 4), (228, 161, 0)] # Green and blue | |
| CLASSES = ['capsules', 'tablets'] | |
| def get_prediction(image): | |
| ''' | |
| Gets image from telebot, make predictions, | |
| counts predicted classes, draws dots on image, | |
| returns dict with counts and labelled image | |
| ''' | |
| # Load a model | |
| model = YOLO('yolov8l.pt') # load an official model | |
| model = YOLO('best.pt') # load a custom model | |
| # image = cv2.imread(image_path) | |
| # Get prediction | |
| prediction = model(image) | |
| # Get predicted classes | |
| predicted_classes = prediction[0].boxes.cls | |
| # Get predicted confidence of each class | |
| prediction_confidences = prediction[0].boxes.conf | |
| # Get polygons | |
| polygons = prediction[0].masks.xy | |
| # Convert polygons to int32 | |
| polygons = [polygon.astype(np.int32) for polygon in polygons] | |
| # Create indices mask that shows what is overlapping polygon has smaller confidence score | |
| indices_mask = remove_overlapping_polygons(polygons, prediction_confidences) | |
| # Create new fixed lists with predicted classes and polygons | |
| fixed_predicted_classes = predicted_classes[np.array(indices_mask, dtype=bool)] | |
| fixed_polygons = [polygons[i] for i in range(len(indices_mask)) if indices_mask[i] == 1] | |
| # fixed_predicted_classes = [predicted_classes[i] for i in range(len(indices_mask)) if indices_mask[i] == 1] | |
| # Get counts of classes | |
| unique, counts = torch.unique(fixed_predicted_classes, return_counts=True) | |
| # Get dicts with counts of classes | |
| count_dict = {CLASSES[int(key)]: value for key, value in zip(unique.tolist(), counts.tolist())} | |
| # # Draw polygons | |
| # for polygon, predicted_class in zip(fixed_polygons, fixed_predicted_classes): | |
| # cv2.polylines(image, [polygon], True, COLORS[int(predicted_class)]) | |
| # Draw dots | |
| for polygon, predicted_class in zip(fixed_polygons, fixed_predicted_classes): | |
| # Find center of polygon | |
| center_coordinates = (np.mean(polygon[:, 0], dtype=np.int32), np.mean(polygon[:, 1], dtype=np.int32)) # x and y respectively | |
| # Draw a circle | |
| cv2.circle(image, center_coordinates, 5, COLORS[int(predicted_class)], 2, cv2.LINE_AA) | |
| # # Show image with predictions on it | |
| # cv2.imshow("Image", image) | |
| # cv2.waitKey(0) | |
| # cv2.destroyAllWindows() | |
| # from google.colab.patches import cv2_imshow | |
| # cv2_imshow(image) | |
| return image, count_dict | |
| def remove_overlapping_polygons(polygons, prediction_confidences): | |
| ''' | |
| Takes polygons, finds overlapping regions, | |
| intersection area, overlap percentage, | |
| creates indices mask that shows what | |
| overlapping polygon has smaller confidence. | |
| ''' | |
| # Convert the NumPy arrays to Shapely polygons | |
| shapely_polygons = [Polygon(polygon) for polygon in polygons] | |
| # Create an empty list with overlapping pairs | |
| overlapping_pairs = [] | |
| # Check for overlaps between all pairs of polygons | |
| for i in range(len(shapely_polygons)): | |
| for j in range(i+1, len(shapely_polygons)): | |
| if shapely_polygons[i].intersects(shapely_polygons[j]): | |
| # Calculate the percentage of overlap | |
| intersection_area = shapely_polygons[i].intersection(shapely_polygons[j]).area | |
| overlap_percentage = intersection_area / shapely_polygons[i].area | |
| # Add overlapping polygons indexes to list | |
| if overlap_percentage > 0.5: | |
| overlapping_pairs.append((i, j)) | |
| # Mask of remains indices | |
| indices_mask = [1 for i in range(len(shapely_polygons))] | |
| # Remove one of the overlapping polygons | |
| for first_over_polygon_ind, second_over_polygon_ind in overlapping_pairs: | |
| # Find index that has the smallest prediction confidence | |
| first_has_bigger_conf = prediction_confidences[first_over_polygon_ind] >= prediction_confidences[second_over_polygon_ind] | |
| index_small_conf = [first_over_polygon_ind, second_over_polygon_ind][first_has_bigger_conf] | |
| # Set value with smaller confidence to 0 in indices_mask | |
| indices_mask[index_small_conf] = 0 | |
| return indices_mask | |