janasumit2911's picture
Update app.py
abb11c6 verified
Raw
History Blame Contribute Delete
3.16 kB
import gradio as gr
from PIL import Image
from ultralytics import YOLO
import requests
import json
import logging
logging.basicConfig(level=logging.INFO)
model = YOLO("Single_Object_BB_Detection_v1.pt")
def detect_objects(images):
results = model(images, max_det=1)
all_bboxes = []
all_bboxes2 = []
for result in results:
boxes = result.boxes.xywhn.tolist()
boxes2 = result.boxes.xywh.tolist()
all_bboxes.append(boxes)
all_bboxes2.append(boxes2)
return all_bboxes, all_bboxes2
def create_solutions(image_urls, all_bboxes, all_bboxes2):
solutions = []
img_id = 1
box_id = 1
category_id = 1
for image_url, bboxes, bboxes2 in zip(image_urls, all_bboxes, all_bboxes2):
for box, box2 in zip(bboxes, bboxes2):
if isinstance(box2[0], list):
w = box2[0][2]
h = box2[0][3]
else:
w = box2[2]
h = box2[3]
area = w * h
seg = [[]]
ans = {"segmentation": seg,"area": area,"iscrowd": 0,"image_id": img_id,"bbox": box,"category_id": category_id,"id": box_id}
obj = {"url": image_url, "answer": [ans]}
solutions.append(obj)
box_id += 1
img_id += 1
return solutions
# def send_results_to_api(data, result_url):
# # Example function to send results to an API
# headers = {"Content-Type": "application/json"}
# response = requests.post(result_url, json=data, headers=headers)
# if response.status_code == 200:
# return response.json() # Return any response from the API if needed
# else:
# return {"error": f"Failed to send results to API: {response.status_code}"}
def process_images(params):
try:
params = json.loads(params)
except json.JSONDecodeError as e:
logging.error(f"Invalid JSON input: {e.msg} at line {e.lineno} column {e.colno}")
return {"error": f"Invalid JSON input: {e.msg} at line {e.lineno} column {e.colno}"}
image_urls = params.get("urls", [])
# api = params.get("api", "")
# job_id = params.get("job_id", "")
if not image_urls:
logging.error("Missing required parameters: 'urls'")
return {"error": "Missing required parameters: 'urls'"}
try:
images = [Image.open(requests.get(url, stream=True).raw) for url in image_urls] # images from URLs
except Exception as e:
logging.error(f"Error loading images: {e}")
return {"error": f"Error loading images: {str(e)}"}
all_bboxes, all_bboxes2 = detect_objects(images) # Perform object detection
solutions = create_solutions(image_urls, all_bboxes, all_bboxes2) # Create solutions with image URLs and bounding boxes
# result_url = f"{api}/{job_id}"
# send_results_to_api(solutions, result_url)
return json.dumps({"solutions": solutions})
inputt = gr.Textbox(label="Parameters (JSON format)")
outputt = gr.JSON()
application = gr.Interface(fn=process_images, inputs=inputt, outputs=outputt, title="Single Object Detection with API Integration")
application.launch()