D3V1L1810's picture
Update app.py
b8f09d0 verified
Raw
History Blame Contribute Delete
3.06 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("Vehicles_Classify_v1.pt")
def detect_objects(images):
results = model(images)
# classes={ 0:"family sedan", 1:"bus", 2:"fire engine",3:"heavy truck", 4:"jeep", 5:"mini bus", 6:"racing car", 7:"SUV", 8:"taxi", 9:"truck" }
classes={ 0:"SUV", 1:"bus", 2:"family sedan", 3:"fire engine",4:"heavy truck", 5:"jeep", 6:"mini bus", 7:"racing car", 8:"taxi", 9:"truck" }
names=[]
for result in results:
probs = result.probs.top1
names.append(classes[probs])
return names
def create_solutions(image_urls, names, file_ids):
solutions = []
for image_url, prediction, file_id in zip(image_urls, names, file_ids):
prediction_list=[]
prediction_list.append(prediction)
obj = {"image": image_url, "answer": prediction_list, "qcUser" : None, "normalfileID" : file_id }
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", [])
if not params.get("normalfileID",[]):
file_ids = [None]*len(image_urls)
else:
file_ids = params.get("normalfileID",[])
# 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, file_ids) # 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="Vehicle Classification with API Integration")
application.launch()