Update app.py

app.py

@@ -1,3 +1,6 @@
+print("import os")
+import os
+
 print("import io")
 import io
 
@@ -13,8 +16,8 @@ import base64
 print("import opencv")
 import cv2
 
-print("import torch")
-import torch
+print("import tempfile")
+import tempfile
 
 print("import gradio")
 import gradio as gr
@@ -35,10 +38,71 @@ print("import scipy")
 from scipy.ndimage import gaussian_filter
 
 print("Initialize Roboflow")
+
+
 rf = Roboflow(api_key="MjzjT2w8u8tlxjmUYDAd")
 project = rf.workspace().project("sku-110k")
 model = project.version(2).model
 
+def encode_image(image_path):
+    """
+    Encode an image located at image_path to base64 string.
+
+    Parameters:
+    - image_path: Path to the image file.
+
+    Returns:
+    A base64 encoded string representing the image.
+    """
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')
+    
+
+def gradio_infer(image, model_name, model_version, confidence, overlap, api_key, output_format, include_labels, stroke_width):
+    '''
+        Send the image to Roboflow API for inference.
+        Returns JSON and image with bounding boxes drawn on to it.
+    '''
+    # Save the incoming image to a temporary file
+    temp_dir = tempfile.mkdtemp()
+    temp_file_path = os.path.join(temp_dir, "temp_image.jpg")
+    image.save(temp_file_path)
+
+    json_url = f"https://detect.roboflow.com/{model_name}/{model_version}?api_key=MjzjT2w8u8tlxjmUYDAd&confidence={confidence}&overlap={overlap}&format=json"
+    image_url = f"https://detect.roboflow.com/{model_name}/{model_version}?api_key=MjzjT2w8u8tlxjmUYDAd&confidence={confidence}&overlap={overlap}&format=image&labels={str(include_labels).lower()}&stroke={stroke_width}"
+    
+    encoded_image = encode_image(temp_file_path)
+    headers = {"Content-Type": "application/x-www-form-urlencoded"}
+
+    json_request = requests.post(json_url, data=encoded_image, headers=headers)
+    image_request = requests.post(image_url, data=encoded_image, headers=headers)
+
+    print("JSON Response Headers:", json_request.headers)
+    print("Image Response Headers:", image_request.headers)
+
+    json_response = {}
+    image_response = None
+
+    if json_request.status_code == 200:
+        try:
+            json_response = json_request.json()
+        except json.JSONDecodeError:
+            json_response = {"error": "Invalid JSON response"}
+    else:
+        json_response = {"error": f"Failed to get JSON response, status code: {json_request.status_code}"}
+
+    if image_request.status_code == 200 and 'image' in image_request.headers.get('Content-Type', ''):
+        try:
+            image_response = Image.open(io.BytesIO(image_request.content))
+        except Exception as e:
+            image_response = None
+            print(f"Failed to open image: {e}")
+    else:
+        print(f"Failed to retrieve image, status code: {image_request.status_code}")
+        print("Image Response Content:", image_request.content)
+
+    return image_response, json_response
+
 def resize_image(image, max_size=1500):
     if max(image.size) > max_size:
         ratio = max_size / float(max(image.size))
@@ -49,17 +113,7 @@ def resize_image(image, max_size=1500):
     buffered = buffer.getvalue()
     return base64.b64encode(buffered).decode("utf-8")
 
-def gradio_infer(image, model="sku-110k", version="2", api_key="gHiUgOSq9GqTnRy5mErk", confidence=0.4, overlap=0.2, format="json", labels=False, stroke=1):
-    base_url = f"https://detect.roboflow.com/{model}/{version}?api_key={api_key}&confidence={confidence}&overlap={overlap}&format={format}"
-    if format == "image" and labels:
-        base_url += "&labels=on"
-    base_url += f"&stroke={stroke}"
-    image_data = resize_image(image)
-    response = requests.post(base_url, data=image_data, headers={"Content-Type": "application/x-www-form-urlencoded"})
-    if format == "json":
-        return response.json()
-    elif format == "image":
-        return Image.open(io.BytesIO(response.content))
+
 
 title = "<center>Cigarette Pack Counter</center>"
 description = "<center><a href='http://counttek.online'><img width='25%' height='25%' src='https://mvp-83056e96f7ab.herokuapp.com/static/countteklogo2.png'></a><br><a href='https://nolenfelten.github.io'>Project by Nolen Felten</a></center>"