Spaces:
Build error
Build error
File size: 5,442 Bytes
5494ea8 |
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 132 133 134 135 136 |
import math
import requests
import numpy as np
import json
from io import BytesIO
from fastapi import FastAPI
from pydantic import BaseModel
from ultralytics import YOLO
from PIL import Image
app = FastAPI()
model = YOLO("posmPJSTRIKE_v1.3.pt")
with open("base_width.json", "r") as f:
base_width = json.load(f)
with open("name_conversion.json", "r") as f:
name_convert = json.load(f)
class ImageRequest(BaseModel):
image_url: str
def get_image_from_url(image_url):
response = requests.get(image_url)
image = Image.open(BytesIO(response.content)).convert("RGB")
return np.array(image)
def name_conversion(actual_distances,object_positions, name_convert):
actual_distances_sys = []
object_positions_sys = {}
for item in actual_distances:
actual_distances_sys.append({'object':(name_convert[list(item.values())[0][0]],name_convert[list(item.values())[0][1]]),'distances': str(list(item.values())[1]) + " cm"})
for item in object_positions:
object_positions_sys.update({name_convert[item]:{"top": str(object_positions[item]['top']) + " cm", "bottom": str(object_positions[item]['bottom']) + " cm", "left": str(object_positions[item]['left']) + " cm", "right": str(object_positions[item]['right']) + " cm"}})
return object_positions_sys, actual_distances_sys
def find_position(objects_names_points, par_pix_cm, image):
object_positions = {}
for obj in objects_names_points:
name = list(obj.keys())[0]
points = list(obj.values())[0]
top_distance = round((points[0][1] - 0) * par_pix_cm[name], 2)
bottom_distance = round((image.size[1] - points[3][1]) * par_pix_cm[name], 2)
left_distance = round((points[0][0] - 0) * par_pix_cm[name], 2)
right_distance = round((image.size[0] - points[3][0]) * par_pix_cm[name], 2)
object_positions.update({name: {"top": top_distance, "bottom": bottom_distance, "left": left_distance, "right": right_distance}})
return object_positions
def get_actual_distance(closest_points, par_pix_cm):
actual_results_n_distance = []
for i in closest_points:
avg_px_cm = ((par_pix_cm[i[0]] + par_pix_cm[i[1]]) / 2)
actual_results_n_distance.append({'object': i, 'distances': round(closest_points[i] * avg_px_cm, 2)})
return actual_results_n_distance
def pixel_per_cm(objects_names_width_pix):
par_pix_cm = {}
for i in objects_names_width_pix:
par_pix_cm_width = base_width[i][0] / objects_names_width_pix[i][0]
par_pix_cm_height = base_width[i][1] / objects_names_width_pix[i][1]
avg_par_pix_cm = (par_pix_cm_width + par_pix_cm_height) / 2
par_pix_cm.update({i: avg_par_pix_cm})
return par_pix_cm
def get_points_n_names(results):
objects_names_points = []
objects_names_width_pix = {}
for box, cls in zip(results[0].boxes.xyxy, results[0].boxes.cls):
x1, y1, x2, y2 = map(int, box)
class_name = results[0].names[int(cls)]
width = x2 - x1
height = y2 - y1
objects_names_points.append({class_name: [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]})
objects_names_width_pix.update({class_name: [width, height]})
return objects_names_points, objects_names_width_pix
def euclidean_distance(point1, point2):
dist_pixels = math.sqrt((point2[0] - point1[0])**2 + (point2[1] - point1[1])**2)
return dist_pixels
def find_closest_points(lst):
closest_points = {}
for i in range(len(lst)):
for j in range(i + 1, len(lst)):
list1 = lst[i]
list2 = lst[j]
min_distance = float('inf')
closest_objects_pair = None
for obj1 in list1.values():
points1 = obj1
for obj2 in list2.values():
points2 = obj2
for point1 in points1:
for point2 in points2:
distance = euclidean_distance(point1, point2)
if distance < min_distance:
min_distance = distance
closest_objects_pair = (list1.keys(), list2.keys())
closest_points.update({(list(closest_objects_pair[0])[0], list(closest_objects_pair[1])[0]): round(min_distance, 2)})
return closest_points
@app.post("/process_image")
def process_image(request: ImageRequest):
image = get_image_from_url(request.image_url)
image = Image.fromarray(image)
size = (640, 640)
image.thumbnail(size)
res = model(image)
objects_names_points, objects_names_width_pix = get_points_n_names(res)
par_pix_cm = pixel_per_cm(objects_names_width_pix)
closest_points = find_closest_points(objects_names_points)
actual_distances = get_actual_distance(closest_points, par_pix_cm)
object_positions = find_position(objects_names_points, par_pix_cm, image)
# Remove the distance between the same object
for item in actual_distances[:]:
if item['object'][0] == item['object'][1]:
actual_distances.remove(item)
# Convert the object names to the system names
object_positions_sys, actual_distances_sys = name_conversion(actual_distances,object_positions, name_convert)
return {
"object_positions": object_positions_sys,
"actual_distances": actual_distances_sys
} |