Update app.py

app.py

@@ -34,13 +34,11 @@ from PIL import Image, ImageDraw
 print("import scipy")
 from scipy.ndimage import gaussian_filter
 
-
 REPO_ID = "thoucentric/Shelf_Objects_Detection_Yolov7_Pytorch"
 FILENAME = "best.pt"
 
 yolov7_custom_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
 
-
 # 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)
@@ -51,7 +49,6 @@ rf = Roboflow(api_key="gHiUgOSq9GqTnRy5mErk")
 project = rf.workspace().project("sku-110k")
 model = project.version(2).model
 
-
 def resize_image(image, max_size=1500):
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
@@ -99,8 +96,6 @@ def roboflow(image, confidence, overlap, stroke_width=1, labels=False):
         "image": image_response
     }
 
-
-    
 # Image Splitting and Merging Functionality
 def split_image(image, tile_size=640, overlap=160):
     img_width, img_height = image.size
@@ -116,8 +111,6 @@ 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():
@@ -134,9 +127,6 @@ def merge_bounding_boxes(results, box):
         adjusted_bboxes.append(adjusted_bbox)
     return adjusted_bboxes
 
-
-
-
 def draw_bounding_boxes(image, bounding_boxes):
     draw = ImageDraw.Draw(image)
     for bbox in bounding_boxes:
@@ -145,7 +135,6 @@ 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:
@@ -182,8 +171,6 @@ 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))
@@ -196,13 +183,9 @@ 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)
@@ -215,8 +198,6 @@ 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)
@@ -237,8 +218,7 @@ 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```python
 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)
@@ -260,8 +240,6 @@ 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)
@@ -294,38 +272,36 @@ 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):
-
     '''
         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, )
+    roboflow_inference = roboflow(image=image_input, confidence=roboflow_confidence_threshold_input, overlap=roboflow_IOU_Threshold_input, stroke_width=roboflow_stroke_width_input, labels=roboflow_labels_input)
 
-    roboflow_image = Image.open(io.BytesIO(roboflow_inference["image"]))
-    roboflow_json = roboflow_inference["json"]
-    
-    return None, None, roboflow_image, roboflow_json
+    if roboflow_inference["image"] is None:
+        raise ValueError("Roboflow API did not return a valid image.")
 
+    roboflow_image = roboflow_inference["image"]
+    roboflow_json = json.dumps(roboflow_inference["json"], indent=4)
+
+    return None, None, 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_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")
+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_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_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")
@@ -333,8 +309,6 @@ 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")
 
@@ -347,7 +321,6 @@ 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 = [
     image_input, 
@@ -359,8 +332,6 @@ inputs = [
     roboflow_stroke_width_input,
 ]
 
-
-
 outputs = [
     yolov7_image_output, 
     yolov7_json_output, 
@@ -368,9 +339,8 @@ outputs = [
     roboflow_json_output,
 ]
 
-
 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>"
+description = "<center><a href='http://counttek.online'><img width="25%" 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(
@@ -384,6 +354,5 @@ interface = gr.Interface(
     allow_flagging="never"
 )
 
-
 # Launch Gradio Interfaces
-interface.launch(debug=True)
+interface.launch(debug=True)