janasumit2911's picture
Update app.py
db06c5b verified
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()