Spaces:
Running
Running
File size: 4,285 Bytes
4bbacde |
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 |
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
|