Spaces:
Sleeping
Sleeping
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("Bottles_Cans_Classify_v1.pt") | |
def detect_objects(images): | |
results = model(images) | |
classes={ 0:"bottle_25cl", 1:"bottle_33cl", 2:"bottle_50cl", 3:"bottle_100cl", 4:"bottle_150cl", 5:"bottle_200cl", 6:"can", 7:"reject" } | |
names=[] | |
for result in results: | |
probs = result.probs.top1 | |
names.append(classes[probs]) | |
return names | |
def create_solutions(image_urls, names): | |
solutions = [] | |
for image_url, prediction in zip(image_urls, names): | |
prediction_list=[] | |
prediction_list.append(prediction) | |
obj = {"url": image_url, "answer": prediction_list} | |
solutions.append(obj) | |
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)}"} | |
names = detect_objects(images) # Perform object detection | |
solutions = create_solutions(image_urls, names) # 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) Eg. {'img_url':['a.jpg','b.jpg']}") | |
outputs = gr.JSON() | |
application = gr.Interface(fn=process_images, inputs=inputt, outputs=outputs, title="Bottles Cans Classification with API Integration") | |
application.launch() |