app.py

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:"Reject", 1:"bottle_100cl", 2:"bottle_150cl", 3:"bottle_200cl", 4:"bottle_25cl", 5:"bottle_33cl", 6:"bottle_50cl", 7:"can"  }
    # 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, file_id):
    solutions = []   
    for image_url, prediction, file in zip(image_urls, names, file_id):       
        prediction_list=[]
        prediction_list.append(prediction)
        obj = {"image": image_url, "answer": prediction_list, "qcUser" : None, "normalfileID": file}      
        solutions.append(obj)   
        print("object prepared")
    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:
        print(f"operating on  : {image_urls}")
        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
    print("prediction done")

    # result_url = f"{api}/{job_id}"
    # send_results_to_api(solutions, result_url)
    print("solution sent...")
    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()