Update app.py

app.py

@@ -34,331 +34,67 @@ 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)
-
-# Roboflow
 print("Initialize Roboflow")
 rf = Roboflow(api_key="MjzjT2w8u8tlxjmUYDAd")
 project = rf.workspace().project("sku-110k")
 model = project.version(2).model
 
-
-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=MjzjT2w8u8tlxjmUYDAd&confidence={confidence}&overlap={overlap}&format=json"
-    image_url = f"https://detect.roboflow.com/sku-110k/2?api_key=MjzjT2w8u8tlxjmUYDAd&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)
-
-    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 {
-        "json": json_response,
-        "image": image_response
-    }
-
-# Image Splitting and Merging Functionality
-def split_image(image, tile_size=640, overlap=160):
-    img_width, img_height = image.size
-    tiles = []
-
-    step = tile_size - overlap
-
-    for top in range(0, img_height, step):
-        for left in range(0, img_width, step):
-            box = (left, top, left + tile_size, top + tile_size)
-            tile = image.crop(box)
-            tiles.append((tile, box))
-
-    return tiles
-
-def merge_bounding_boxes(results, box):
-    adjusted_bboxes = []
-    for idx, row in results.pandas().xyxy[0].iterrows():
-        adjusted_bbox = {
-            "class": row['name'],
-            "center_x": (row['xmin'] + row['xmax']) / 2 + box[0],
-            "center_y": (row['ymin'] + row['ymax']) / 2 + box[1],
-            "xmin": row['xmin'] + box[0],
-            "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:
-        color = "red"
-        draw.rectangle([bbox['xmin'], bbox['ymin'], bbox['xmax'], bbox['ymax']], outline=color, width=3)
-        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:
-        return []
-
-    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
-
-    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))
-
-    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)
-
-# Edge Enhancement Functions
-def apply_edge_enhancement(image):
-    gray_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-
-    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)
-
-    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)
-    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:
-        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, final_bounding_boxes)
-
-    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)
-
-    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
-
-    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
-
-
-
-
-# Function to resize and encode an image
 def resize_image(image, max_size=1500):
-    print("resize_image initialized")
     if max(image.size) > max_size:
-
-        print("Image size is over the max size, resizing")
         ratio = max_size / float(max(image.size))
-        print("Ratio is: " + str(ratio))
-        
         new_size = tuple([int(x * ratio) for x in image.size])
-        print("New size is: " + str(new_size))
-        
-        print("Resizing image")
         image = image.resize(new_size, Image.LANCZOS)
-
-
-
-    print("Create io.BytesIO instance")
     buffer = io.BytesIO()
-    
-    print("Save image to BytesIO")
     image.save(buffer, format="PNG")
     buffered = buffer.getvalue()
-
-    print("Return base64 encoded image")
     return base64.b64encode(buffered).decode("utf-8")
 
-
-def infer(image, model, version, api_key, confidence=0.4, overlap=0.20, format="json", labels=False, stroke=1):
-    
+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":
-        if labels:
-            base_url += "&labels=on"
-        base_url += f"&stroke={stroke}"
-    
-    print(base_url)
-
-
+    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":
-        print(json.dumps(response.json(), indent=4))
-        return json.dumps(response.json(), indent=4)
+        return response.json()
     elif format == "image":
         return Image.open(io.BytesIO(response.content))
 
-
-def gradio_infer(image, model, version, api_key, confidence, overlap, format, labels, stroke):
-    result_json = infer(image, model, version, api_key, confidence, overlap, "json", labels, stroke)
-    result_image = infer(image, model, version, api_key, confidence, overlap, "image", labels, stroke)
-    
-    return result_image, result_json
-
-
-
-
 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>"
 footer = ("<center><b>Item Classes it will detect (Total 140 Classes)</b></center>")
 
-interface = gr.Interface(
-    fn=gradio_infer,
-    inputs=[
-        gr.Image(type="pil", label="Input Image"),
-        gr.Textbox(value="sku-110k", label="Model Name"),
-        gr.Textbox(value="2", label="Model Version"),
-        gr.Textbox(value="gHiUgOSq9GqTnRy5mErk", label="API Key"),
-        gr.Slider(0.0, 1.0, value=0.40, label="Confidence Threshold"),
-        gr.Slider(0.0, 1.0, value=0.20, label="Overlap Threshold"),
-        gr.Radio(["json", "image"], value="image", label="Output Format"),
-        gr.Checkbox(False, label="Include Labels"),
-        gr.Slider(1, 10, value=5, step=1, label="Stroke Width"),
-    ],
-    outputs=[
-        gr.Image(label="Roboflow Output Image"),
-        gr.Textbox(label="Roboflow JSON Result"),
-    ],
-    title=title,
-    description=description,
-    article=footer,
-    cache_examples=False,
-    allow_flagging="never"
-)
+def create_interface(shelf_number):
+    return gr.Interface(
+        fn=gradio_infer,
+        inputs=[
+            gr.Image(type="pil", label=f"Shelf {shelf_number}"),
+            gr.Textbox(value="sku-110k", label=f"Model Name - Shelf {shelf_number}"),
+            gr.Textbox(value="2", label="Model Version"),
+            gr.Slider(0.0, 1.0, value=0.40, label="Confidence Threshold"),
+            gr.Slider(0.0, 1.0, value=0.20, label="Overlap Threshold"),
+            gr.Textbox(value="gHiUgOSq9GqTnRy5mErk", label="API Key"),
+            gr.Radio(["json", "image"], value="image", label="Output Format"),
+            gr.Checkbox(False, label="Include Labels"),
+            gr.Slider(1, 10, value=5, step=1, label="Stroke Width")
+        ],
+        outputs=[
+            gr.Image(label=f"Roboflow Output Image - Shelf {shelf_number}"),
+            gr.JSON(label=f"Roboflow JSON Result - Shelf {shelf_number}")
+        ],
+        title=title,
+        description=description,
+        article=footer,
+        cache_examples=False,
+        allow_flagging="never"
+    )
+
+# Create a list of interfaces
+interfaces = [create_interface(i) for i in range(1, 12)]
+
+# Combine all interfaces into one layout
+interface = gr.TabbedInterface(interfaces, tab_names=[f"Shelf {i}" for i in range(1, 12)])
 
-# Launch Gradio Interfaces
+# Launch the combined interface
 interface.launch(debug=True)