main.py

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
    }