Update app.py

app.py

@@ -4,15 +4,19 @@ from huggingface_hub import hf_hub_download
 from PIL import Image, ImageDraw
 import numpy as np
 import json
+import cv2
+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)
 
-model = torch.hub.load('Owaiskhan9654/yolov7-1:main', model='custom', path_or_model=yolov7_custom_weights, force_reload=True)  # Github repository https://github.com/Owaiskhan9654
+# Load YOLOv7 Custom Model
+model = torch.hub.load('Owaiskhan9654/yolov7-1:main', model='custom', path_or_model=yolov7_custom_weights, force_reload=True)
 
-def split_image(image, tile_size=640, overlap=80):
+# Image Splitting and Merging Functionality
+def split_image(image, tile_size=640, overlap=160):
     img_width, img_height = image.size
     tiles = []
 
@@ -50,39 +54,76 @@ def draw_bounding_boxes(image, bounding_boxes):
         draw.text((bbox['xmin'], bbox['ymin']), f"{bbox['class']} {bbox['confidence']:.2f}", fill=color)
     return image
 
-def non_max_suppression(bounding_boxes, iou_threshold):
+# Non-Max Suppression Implementations
+def soft_nms(bounding_boxes, iou_threshold=0.3, sigma=0.5, score_threshold=0.001):
     if not bounding_boxes:
         return []
 
-    bboxes = sorted(bounding_boxes, key=lambda x: x['confidence'], reverse=True)
-    final_bboxes = []
+    def iou(box1, box2):
+        x1 = max(box1['xmin'], box2['xmin'])
+        y1 = max(box1['ymin'], box2['ymin'])
+        x2 = min(box1['xmax'], box2['xmax'])
+        y2 = min(box1['ymax'], box2['ymax'])
+
+        intersection = max(0, x2 - x1) * max(0, y2 - y1)
+        box1_area = (box1['xmax'] - box1['xmin']) * (box1['ymax'] - box1['ymin'])
+        box2_area = (box2['xmax'] - box2['xmin']) * (box2['ymax'] - box2['ymin'])
+
+        union = box1_area + box2_area - intersection
+        return intersection / union
+
+    boxes = sorted(bounding_boxes, key=lambda x: x['confidence'], reverse=True)
+    final_boxes = []
+
+    while boxes:
+        chosen_box = boxes.pop(0)
+        final_boxes.append(chosen_box)
+
+        new_boxes = []
+        for box in boxes:
+            iou_score = iou(chosen_box, box)
+            if iou_score > iou_threshold:
+                box['confidence'] *= np.exp(-(iou_score ** 2) / sigma)
+            if box['confidence'] > score_threshold:
+                new_boxes.append(box)
+        boxes = new_boxes
 
-    while bboxes:
-        chosen_box = bboxes.pop(0)
-        final_bboxes.append(chosen_box)
+    return final_boxes
 
-        def iou(box1, box2):
-            x1 = max(box1['xmin'], box2['xmin'])
-            y1 = max(box1['ymin'], box2['ymin'])
-            x2 = min(box1['xmax'], box2['xmax'])
-            y2 = min(box1['ymax'], box2['ymax'])
+# Density Map Generation and Counting Functions
+def generate_density_map(image, bounding_boxes, sigma=4):
+    density_map = np.zeros((image.height, image.width))
 
-            intersection = max(0, x2 - x1) * max(0, y2 - y1)
-            box1_area = (box1['xmax'] - box1['xmin']) * (box1['ymax'] - box1['ymin'])
-            box2_area = (box2['xmax'] - box2['xmin']) * (box2['ymax'] - box2['ymin'])
+    for bbox in bounding_boxes:
+        center_x = int((bbox['xmin'] + bbox['xmax']) / 2)
+        center_y = int((bbox['ymin'] + bbox['ymax']) / 2)
+        density_map[center_y, center_x] += 1
+
+    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)
 
-            union = box1_area + box2_area - intersection
-            return intersection / union
+# Edge Enhancement Functions
+def apply_edge_enhancement(image):
+    gray_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
 
-        bboxes = [box for box in bboxes if iou(chosen_box, box) < iou_threshold]
+    sobel_x = cv2.Sobel(gray_image, cv2.CV_64F, 1, 0, ksize=3)
+    sobel_y = cv2.Sobel(gray_image, cv2.CV_64F, 0, 1, ksize=3)
+    sobel_combined = np.sqrt(sobel_x**2 + sobel_y**2)
+    sobel_combined = np.uint8(sobel_combined / sobel_combined.max() * 255)
 
-    return final_bboxes
+    enhanced_image = cv2.cvtColor(sobel_combined, cv2.COLOR_GRAY2RGB)
+    return Image.fromarray(enhanced_image)
 
-def object_detection(image, model_choice, conf_threshold=0.25, iou_threshold=0.45):
-    image = Image.fromarray(image)  # Convert to PIL Image
-    model.conf = conf_threshold  # set confidence threshold
-    model.iou = iou_threshold    # set IOU threshold
-    tiles = split_image(image)
+# Object Detection Functions
+def object_detection(image, conf_threshold=0.25, iou_threshold=0.45):
+    image = Image.fromarray(image)
+    model.conf = conf_threshold
+    model.iou = iou_threshold
+
+    tiles = split_image(image, tile_size=640, overlap=160)
     all_bounding_boxes = []
 
     for tile, box in tiles:
@@ -90,49 +131,126 @@ def object_detection(image, model_choice, conf_threshold=0.25, iou_threshold=0.4
         adjusted_bboxes = merge_bounding_boxes(results, box)
         all_bounding_boxes.extend(adjusted_bboxes)
 
-    final_bounding_boxes = non_max_suppression(all_bounding_boxes, iou_threshold)
+    final_bounding_boxes = soft_nms(all_bounding_boxes, iou_threshold=iou_threshold)
     image_with_boxes = draw_bounding_boxes(image, final_bounding_boxes)
 
     json_response = json.dumps(final_bounding_boxes, indent=4)
     return image_with_boxes, json_response
 
-def save_results(image, json_data):
-    image.save("output_image.jpg")
-    with open("output_data.json", "w") as f:
-        f.write(json_data)
-    return "Results saved successfully."
+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)
+
+    model.conf = conf_threshold
+    model.iou = iou_threshold
+
+    tiles = split_image(image_enhanced, tile_size=640, overlap=160)
+    all_bounding_boxes = []
+
+    for tile, box in tiles:
+        results = model(tile)
+        adjusted_bboxes = merge_bounding_boxes(results, box)
+        all_bounding_boxes.extend(adjusted_bboxes)
+
+    final_bounding_boxes = soft_nms(all_bounding_boxes, iou_threshold=iou_threshold)
+    image_with_boxes = draw_bounding_boxes(image_enhanced, final_bounding_boxes)
+
+    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)
+    image_enhanced = apply_edge_enhancement(image)
+
+    model.conf = conf_threshold
+    model.iou = iou_threshold
 
-title = "<center>Yolov7 Custom</center>"
+    tiles = split_image(image_enhanced, tile_size=640, overlap=160)
+    all_bounding_boxes = []
+
+    for tile, box in tiles:
+        results = model(tile)
+        adjusted_bboxes = merge_bounding_boxes(results, box)
+        all_bounding_boxes.extend(adjusted_bboxes)
 
+    final_bounding_boxes = soft_nms(all_bounding_boxes, iou_threshold=iou_threshold)
+    density_map = generate_density_map(image_enhanced, final_bounding_boxes)
+
+    density_map_rescaled = cv2.applyColorMap(
+        (density_map / np.max(density_map) * 255).astype(np.uint8),
+        cv2.COLORMAP_JET
+    )
+    density_map_pil = Image.fromarray(cv2.cvtColor(density_map_rescaled, cv2.COLOR_BGR2RGB))
+    image_with_density = Image.blend(image_enhanced, density_map_pil, alpha=0.5)
+
+    json_response = json.dumps(final_bounding_boxes, indent=4)
+    object_count = count_from_density_map(density_map)
+
+    summary = json.dumps({"object_count": int(object_count)}, indent=4)
+    return image_with_density, json_response, summary
+
+# Gradio Interface Definitions
 inputs = [
     gr.inputs.Image(shape=(4080, 1836), image_mode="RGB", source="upload", label="Upload Image", optional=False),
-    gr.inputs.Dropdown(["best.pt"], default="best.pt", label="Model"),
     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_image = gr.outputs.Image(type="pil", label="Output Image")
 outputs_json = gr.Textbox(label="Bounding Boxes JSON")
 
-Custom_description = "<center>Nolen Felten</center>"
-
-Footer = ("<br><br><br><br><center><b>Item Classes it will detect(Total 140 Classes) <br></b>")
-examples1 = [["Images/Image1.jpg"]]
-css = ".output-image {height: 50rem important; width: 100% !important;}, .input-image {height: 50rem !important; width: 100% !important;}"
-css = ".image-preview {height: auto important;}"
+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>")
 
+# Regular Object Detection Interface
 interface = gr.Interface(
     fn=object_detection,
     inputs=inputs,
     outputs=[outputs_image, outputs_json],
-    title="Cig Counter",
-    description=Custom_description,
-    article=Footer,
+    title=title,
+    description=description,
+    article=footer,
+    cache_examples=False,
+    allow_flagging="never"
+)
+
+# Edge Enhanced Object Detection Interface
+interface_edge = 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_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',
-    examples=examples1
+    allow_flagging="never"
 )
 
-interface.launch(debug=True)
+# Launch Gradio Interfaces
+def launch_interfaces():
+    interface.launch(debug=True)
+    interface_edge.launch(debug=True)
+    interface_density_edge.launch(debug=True)
 
-save_button = gr.Button("Save Results")
-save_button.click(fn=save_results, inputs=[outputs_image, outputs_json], outputs="text")
+if __name__ == "__main__":
+    launch_interfaces()