Update app.py

app.py

@@ -1,109 +1,91 @@
 import gradio as gr
 from PIL import Image
 from ultralytics import YOLO
-import requests
+import asyncio
+import aiohttp
 import json
 import logging
 import cv2
-from numpy import asarray
+from io import BytesIO
 
 logging.basicConfig(level=logging.INFO)
 
 model_detection = YOLO('./detection_best.pt')
 model_classification = YOLO('./classification_best.pt')
 
-def detect_objects(images):
-    classes={ 2: "Positive", 1: "Negative"}
-    names = [] 
-    for image in images:
-        image = asarray(image)
-        image = cv2.resize(image,(640,640))
-        results_detection = model_detection(image)
-        print(results_detection)
-    # Load the image using OpenCV
-        # img = cv2.imread(image_path)
-        # Process each detected object
-        if results_detection:
-            for result in results_detection:
-                for box in result.boxes:
-                    # Get bounding box coordinates (x1, y1, x2, y2)
-                    x1, y1, x2, y2 = map(int, box.xyxy[0])
-            
-                    # Crop the bounding box from the image
-                    cropped_img = image[y1:y2, x1:x2]
-            
-                    # Resize the cropped image to 640x640
-                    resized_img = cv2.resize(cropped_img, (640, 640))
-                    resized_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)
-                    # Perform inference using the classification model
-                    results_classification = model_classification.predict(resized_img,save = True)
-            
-                    detected = False
-                    print(results_classification)
-                    # Process classification results
-                    for res in results_classification:
-                        # Get class probabilities and labels
-                        top1_class = res.probs.top1  # Predicted class
-                        top1_confidence = res.probs.top1conf
-                        print(top1_class)
-                        names.append([classes[top1_class]])
-        else:
-            names.append(['None'])
-        # print(names)
-    return names
+async def fetch_image(url):
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url) as response:
+            if response.status == 200:
+                image_data = await response.read()
+                return Image.open(BytesIO(image_data))
+            else:
+                logging.error(f"Failed to load image from {url}")
+                return None
+
+async def detect_objects(images):
+    classes = {2: "Positive", 1: "Negative"}
+    results = []
+
+    processed_images = [cv2.resize(np.array(image), (640, 640)) for image in images]
+
+    results_detection = model_detection(processed_images)
+    
+    for image, detection in zip(processed_images, results_detection):
+        names = []
+        if detection:
+            for box in detection.boxes:
+                x1, y1, x2, y2 = map(int, box.xyxy[0])
+                cropped_img = image[y1:y2, x1:x2]
+                resized_img = cv2.resize(cropped_img, (640, 640))
+                resized_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)
+
+                results_classification = model_classification.predict(resized_img)
+                
+                if results_classification:
+                    top1_class = results_classification[0].probs.top1
+                    names.append(classes[top1_class])
+        if not names:
+            names.append("None")
+        results.append(names)
+    return results
 
 def create_solutions(image_urls, names, file_ids):
-    solutions = []   #list to store all the objects
-
-    for image_url, class_name, file_id in zip(image_urls, names, file_ids):      
-        obj = {"image": image_url, "answer": [class_name], "qcUserId" : 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}"}
+    return [
+        {"image": url, "answer": name, "qcUserId": None, "normalfileID": file_id}
+        for url, name, file_id in zip(image_urls, names, file_ids)
+    ]
 
-def process_images(params):
+async def process_images_async(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}"}
-    
+        logging.error(f"Invalid JSON input: {e}")
+        return {"error": f"Invalid JSON input: {e}"}
+
     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", "")
-    
+    file_ids = params.get("normalfileID", [None] * len(image_urls))
+
     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
+    images = await asyncio.gather(*[fetch_image(url) for url in image_urls])
 
-    # result_url = f"{api}/{job_id}"
-    # send_results_to_api(solutions, result_url)
+    if not any(images):
+        logging.error("No valid images were loaded.")
+        return {"error": "No valid images were loaded."}
+
+    names = await detect_objects(images)
+    solutions = create_solutions(image_urls, names, file_ids)
 
     return json.dumps({"solutions": solutions})
 
+def process_images(params):
+    return asyncio.run(process_images_async(params))
+
 inputt = gr.Textbox(label="Parameters (JSON format) Eg. img_url:['','']")
 outputs = gr.JSON()
 
 application = gr.Interface(fn=process_images, inputs=inputt, outputs=outputs, title="ART +ve -ve Detection with API Integration")
-application.launch()
+application.launch()