Update app.py

app.py

@@ -1,37 +1,85 @@
+print("import io")
+import io
+
+print("import requests")
+import requests
+
+print("import json")
+import json
+
+print("import base64")
+import base64
+
+print("import opencv")
+import cv2
+
 print("import torch")
 import torch
 
 print("import gradio")
 import gradio as gr
 
+print("import numpy")
+import numpy as np
+
+print("import Roboflow")
+from roboflow import Roboflow
+
 print("import huggingface_hub")
 from huggingface_hub import hf_hub_download
 
 print("import PIL")
 from PIL import Image, ImageDraw
 
-print("import numpy")
-import numpy as np
-
-print("import json")
-import json
-
-print("import opencv")
-import cv2
-
 print("import scipy")
 from scipy.ndimage import gaussian_filter
 
 
 # Constants and Model Downloads
-REPO_ID = "thoucentric/Shelf_Objects_Detection_Yolov7_Pytorch"
-FILENAME = "best.pt"
-yolov7_custom_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+print("Download YOLO")
+yolov7_custom_weights = hf_hub_download(repo_id = "nolenfelten/Shelf_Objects_Detection_Yolov7_Pytorch", filename = "best.pt")
 
 # Load YOLOv7 Custom Model
 print("Load YOLOv7 Custom Model")
 model = torch.hub.load('Owaiskhan9654/yolov7-1:main', model='custom', path_or_model=yolov7_custom_weights, force_reload=True)
 
+# Roboflow
+print("Initialize Roboflow")
+rf = Roboflow(api_key="gHiUgOSq9GqTnRy5mErk")
+project = rf.workspace().project("sku-110k")
+model = project.version(2).model
+
+
+def encode_image(image):
+    buffered = io.BytesIO()
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+
+
+def roboflow(image, confidence, overlap, stroke_width=1, labels=False):
+    '''
+        Send the image to Roboflow API for inference.
+        Returns JSON and image with bounding boxes drawn on to it.
+    '''
+
+    json_url = f"https://detect.roboflow.com/sku-110k/2?api_key=gHiUgOSq9GqTnRy5mErk&confidence={confidence}&overlap={overlap}&format=json"
+    image_url = f"https://detect.roboflow.com/sku-110k/2?api_key=gHiUgOSq9GqTnRy5mErk&confidence={confidence}&overlap={overlap}&format=image&labels={str(labels).lower()}&stroke={stroke_width}"
+    
+    encoded_image = encode_image(image)
+    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)
+
+    response = {
+        "json": json_request.content,
+        "image": image_request.content
+    }
+
+    return response
+        
+
+    
 # Image Splitting and Merging Functionality
 def split_image(image, tile_size=640, overlap=160):
     img_width, img_height = image.size
@@ -47,6 +95,8 @@ def split_image(image, tile_size=640, overlap=160):
 
     return tiles
 
+
+
 def merge_bounding_boxes(results, box):
     adjusted_bboxes = []
     for idx, row in results.pandas().xyxy[0].iterrows():
@@ -55,14 +105,17 @@ def merge_bounding_boxes(results, box):
             "center_x": (row['xmin'] + row['xmax']) / 2 + box[0],
             "center_y": (row['ymin'] + row['ymax']) / 2 + box[1],
             "xmin": row['xmin'] + box[0],
-            "ymin": row['ymin'] + box[1],
             "xmax": row['xmax'] + box[0],
+            "ymin": row['ymin'] + box[1],
             "ymax": row['ymax'] + box[1],
             "confidence": row['confidence'],
         }
         adjusted_bboxes.append(adjusted_bbox)
     return adjusted_bboxes
 
+
+
+
 def draw_bounding_boxes(image, bounding_boxes):
     draw = ImageDraw.Draw(image)
     for bbox in bounding_boxes:
@@ -71,6 +124,7 @@ def draw_bounding_boxes(image, bounding_boxes):
         draw.text((bbox['xmin'], bbox['ymin']), f"{bbox['class']} {bbox['confidence']:.2f}", fill=color)
     return image
 
+
 # Non-Max Suppression Implementations
 def soft_nms(bounding_boxes, iou_threshold=0.3, sigma=0.5, score_threshold=0.001):
     if not bounding_boxes:
@@ -107,6 +161,8 @@ def soft_nms(bounding_boxes, iou_threshold=0.3, sigma=0.5, score_threshold=0.001
 
     return final_boxes
 
+
+
 # Density Map Generation and Counting Functions
 def generate_density_map(image, bounding_boxes, sigma=4):
     density_map = np.zeros((image.height, image.width))
@@ -119,9 +175,13 @@ def generate_density_map(image, bounding_boxes, sigma=4):
     density_map = gaussian_filter(density_map, sigma=sigma)
     return density_map
 
+
+
 def count_from_density_map(density_map, threshold=0.05):
     return np.sum(density_map > threshold)
 
+
+
 # Edge Enhancement Functions
 def apply_edge_enhancement(image):
     gray_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
@@ -134,6 +194,8 @@ def apply_edge_enhancement(image):
     enhanced_image = cv2.cvtColor(sobel_combined, cv2.COLOR_GRAY2RGB)
     return Image.fromarray(enhanced_image)
 
+
+
 # Object Detection Functions
 def object_detection(image, conf_threshold=0.25, iou_threshold=0.45):
     image = Image.fromarray(image)
@@ -154,6 +216,8 @@ def object_detection(image, conf_threshold=0.25, iou_threshold=0.45):
     json_response = json.dumps(final_bounding_boxes, indent=4)
     return image_with_boxes, json_response
 
+
+
 def object_detection_with_edge_enhancement(image, conf_threshold=0.25, iou_threshold=0.45):
     image = Image.fromarray(image)
     image_enhanced = apply_edge_enhancement(image)
@@ -175,6 +239,8 @@ def object_detection_with_edge_enhancement(image, conf_threshold=0.25, iou_thres
     json_response = json.dumps(final_bounding_boxes, indent=4)
     return image_with_boxes, json_response
 
+
+
 def object_detection_density_edge(image, conf_threshold=0.25, iou_threshold=0.45):
     """Apply edge enhancement and density-based counting."""
     image = Image.fromarray(image)
@@ -207,73 +273,101 @@ def object_detection_density_edge(image, conf_threshold=0.25, iou_threshold=0.45
     summary = json.dumps({"object_count": int(object_count)}, indent=4)
     return image_with_density, json_response, summary
 
+
+
+def procedure(image_input, yolov7_confidence_threshold_input, yolov7_IOU_Threshold_input, roboflow_confidence_threshold_input, roboflow_IOU_Threshold_input, roboflow_labels_input, roboflow_stroke_width_input, yolov7_image_output, yolov7_json_output, roboflow_image_output, roboflow_json_output):
+    '''
+        This function takes in an image and applies both YOLOv7 and Roboflow object detection models to it. 
+        It then returns the images and JSON results.
+    '''
+
+    print("Begin Roboflow inferences.")
+    roboflow_inference = roboflow(image = image_input, labels=roboflow_labels_input, stroke_width=roboflow_stroke_width_input, confidence = roboflow_confidence_threshold_input, overlap = roboflow_IOU_Threshold_input, )
+
+    yolov7_image, yolov7_json = object_detection(np.array(image_input), yolov7_confidence_threshold_input, yolov7_IOU_Threshold_input)
+    
+    roboflow_image = Image.open(io.BytesIO(roboflow_inference["image"]))
+    roboflow_json = roboflow_inference["json"]
+
+    return yolov7_image, yolov7_json, roboflow_image, roboflow_json
+
+
+# Uploaded image.
+image_input = gr.Image(shape=(4080, 1836), image_mode="RGB", source="upload", label="Upload Image", optional=False)
+
+# YOLOv7 Confidence Threshold input.
+yolov7_confidence_threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value = 0.45, step=0.01, label="YOLOv7 Confidence Threshold")
+
+# YOLOv7 IOU Threshold.
+yolov7_IOU_Threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value = 0.45, step=0.01, label="YOLOv7 IOU Threshold")
+
+# Roboflow Confidence Threshold input.
+roboflow_confidence_threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value = 0.45, step=0.01, label="Roboflow Confidence Threshold")
+
+# Roboflow IOU Threshold.
+roboflow_IOU_Threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value = 0.45, step=0.01, label="Roboflow IOU Threshold")
+
+# Roboflow Labels.
+roboflow_labels_input = gr.Checkbox(label="Roboflow Labels")
+
+# Roboflow Stroke Width.
+roboflow_stroke_width_input = gr.Radio([1, 2, 5, 10], label="Stroke Width")
+
+
+
+# YOLOv7 Image Output.
+yolov7_image_output = gr.Image(type="pil", label="YOLOv7 Output Image")
+
+# YOLOv7 JSON Output.
+yolov7_json_output = gr.Textbox(label="YOLOv7 Bounding Boxes JSON")
+
+# Roboflow Image Output.
+roboflow_image_output = gr.Image(type="pil", label="Roboflow Output Image")
+
+# Roboflow JSON Output.
+roboflow_json_output = gr.Textbox(label="Roboflow Bounding Boxes JSON")
+
+
 # Gradio Interface Definitions
 inputs = [
-    gr.inputs.Image(shape=(4080, 1836), image_mode="RGB", source="upload", label="Upload Image", optional=False),
-    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.01, label="Confidence Threshold"),
-    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.01, label="IOU Threshold"),
+    image_input, 
+    yolov7_confidence_threshold_input, 
+    yolov7_IOU_Threshold_input, 
+    roboflow_confidence_threshold_input,
+    roboflow_IOU_Threshold_input,
+    roboflow_labels_input,
+    roboflow_stroke_width_input,
 ]
-outputs_image = [
-    gr.outputs.Image(type="pil", label="Output Image")
-]
-outputs_json = gr.Textbox(label="Bounding Boxes JSON")
-
-title = "<center>Yolov7 Custom Object Detection</center>"
-description = "<center>Nolen Felten</center>"
-footer = ("<br><br><center><b>Item Classes it will detect (Total 140 Classes)</b></center>")
-
-
-interfaces = [
-    # Regular Object Detection Interface
-    gr.Interface(
-        fn=object_detection,
-        inputs=inputs,
-        outputs=[outputs_image, outputs_json],
-        title=title,
-        description=description,
-        article=footer,
-        cache_examples=False,
-        allow_flagging="never"
-    ), 
-    
 
-    # Edge Enhanced Object Detection Interface
-    gr.Interface(
-        fn=object_detection_with_edge_enhancement,
-        inputs=inputs,
-        outputs=[outputs_image, outputs_json],
-        title="Object Detection with Edge Enhancement",
-        description=description,
-        article=footer,
-        cache_examples=False,
-        allow_flagging="never"
-    )
-]
 
-# Edge Enhanced Density-Based Counting Interface
-inputs_density_edge = [
-    gr.inputs.Image(shape=(4080, 1836), image_mode="RGB", source="upload", label="Upload Image", optional=False),
-    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.01, label="Confidence Threshold"),
-    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.01, label="IOU Threshold"),
+
+outputs = [
+    yolov7_image_output, 
+    yolov7_json_output, 
+    roboflow_image_output, 
+    roboflow_json_output,
 ]
-outputs_density_edge = [outputs_image, outputs_json, gr.Textbox(label="Density-based Count Summary")]
-
-interface_density_edge = gr.Interface(
-    fn=object_detection_density_edge,
-    inputs=inputs_density_edge,
-    outputs=outputs_density_edge,
-    title="Object Detection with Density-Based Counting and Edge Enhancement",
-    description=description,
-    article=footer,
-    cache_examples=False,
-    allow_flagging="never"
+
+
+title = "<center>Cigarette Pack Counter</center>"
+description = "<center><a href='http://counttek.online'><img src='https://mvp-83056e96f7ab.herokuapp.com/static/countteklogo2.png'></a><br><a href='https://nolenfelten.github.io'>Project by Nolen Felten</a></center>"
+footer = ("<center><b>Item Classes it will detect (Total 140 Classes)</b></center>")
+
+
+
+interface = gr.Interface(
+    # Run this function when user clicks "Submit".
+    fn = procedure,
+
+    inputs,
+
+    outputs,
+    title = title,
+    description = description,
+    article = footer,
+    cache_examples = False,
+    allow_flagging = "never"
 )
 
 # Launch Gradio Interfaces
-def launch_interfaces():
-    interface.launch(debug=True)
-    interface_edge.launch(debug=True)
-    interface_density_edge.launch(debug=True)
-
-if __name__ == "__main__":
-    launch_interfaces()
+interface.launch(debug=True)