Update app.py

app.py

@@ -1,103 +1,125 @@
 import gradio as gr
-import os
+from roboflow import Roboflow
 import tempfile
-import math
-import cv2
+import os
+from sahi.slicing import slice_image
 import numpy as np
-import supervision as sv
-from roboflow import Roboflow
+import cv2
 
-# Initialize Roboflow
+# Inisialisasi Roboflow (for model path)
 rf = Roboflow(api_key="Otg64Ra6wNOgDyjuhMYU")
 project = rf.workspace("alat-pelindung-diri").project("nescafe-4base")
 model = project.version(16).model
 
-# Helper function for SAHI (Supervision Slicing)
-def calculate_tile_size(image_shape: tuple[int, int], tiles: tuple[int, int], overlap_ratio_wh: tuple[float, float] = (0.0, 0.0)):
-    w, h = image_shape
-    rows, columns = tiles
-    tile_width = math.ceil(w / columns * (1 + overlap_ratio_wh[0]))
-    tile_height = math.ceil(h / rows * (1 + overlap_ratio_wh[1]))
-    overlap_wh = (math.ceil(tile_width * overlap_ratio_wh[0]), math.ceil(tile_height * overlap_ratio_wh[1]))
-    return (tile_width, tile_height), overlap_wh
-
-# Function to handle inference and tiles
-def detect_objects(image):
-    # Convert PIL image to NumPy array (for OpenCV compatibility)
-    img = np.array(image)  # Gradio image is in PIL format, convert it to NumPy array
-    img_rgb = img  # Keep the image as RGB format, avoid unnecessary conversion to BGR
+# Fungsi untuk melakukan Non-Maximum Suppression (NMS)
+def apply_nms(predictions, iou_threshold=0.5):
+    boxes = []
+    scores = []
+    classes = []
 
-    image_shape = (img.shape[1], img.shape[0])
+    # Extract boxes, scores, and class info
+    for prediction in predictions:
+        boxes.append(prediction['bbox'])
+        scores.append(prediction['confidence'])
+        classes.append(prediction['class'])
 
-    # Parameters for slicing (tiles and overlap)
-    tiles = (8, 8)  # Use 8x8 tiles for better detection of small objects
-    overlap_ratio_wh = (0.2, 0.2)  # 20% overlap between tiles for better context
-    slice_wh, overlap_wh = calculate_tile_size(image_shape, tiles, overlap_ratio_wh)
+    boxes = np.array(boxes)
+    scores = np.array(scores)
+    classes = np.array(classes)
 
-    # Generate offsets but don't visualize the tiles with rectangles (remove the drawing step)
-    offsets = sv.InferenceSlicer._generate_offset(image_shape, slice_wh, None, overlap_wh)
-    tiled_image = img_rgb.copy()
+    # Perform NMS using OpenCV
+    indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), score_threshold=0.25, nms_threshold=iou_threshold)
 
-    # Save the PIL image to a temporary file for Roboflow model prediction
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
-        image.save(temp_file, format="JPEG")
-        temp_file_path = temp_file.name
+    # Convert tuple of indices to a flat NumPy array
+    indices = indices.flatten() if isinstance(indices, tuple) else indices
 
-    # Annotate with Roboflow model predictions using the temporary file path
-    predictions = model.predict(temp_file_path, confidence=40, overlap=30).json()  # Adjusted confidence for small object detection
-    class_count = {}
+    nms_predictions = []
 
-    # Define a color palette for different classes
-    color_palette = {
-        "bearbrand": (0, 255, 0),   # Green for class 1
-        "nescafe latte": (0, 0, 255),   # Red for class 2
-        "nescafe original": (255, 0, 0),   # Blue for class 3
-        "nescafe mocha": (0, 255, 255) # Yellow for class 4
-        #"class_5": (255, 0, 255)  # Magenta for class 5
-        # You can add more colors based on the number of classes you have
-    }
+    for i in indices:
+        nms_predictions.append({
+            'class': classes[i],
+            'bbox': boxes[i],
+            'confidence': scores[i]
+        })
 
-    # Draw bounding boxes with different colors and label classes
-    for prediction in predictions['predictions']:
-        x1 = int(prediction['x'] - prediction['width'] / 2)
-        y1 = int(prediction['y'] - prediction['height'] / 2)
-        x2 = int(prediction['x'] + prediction['width'] / 2)
-        y2 = int(prediction['y'] + prediction['height'] / 2)
+    return nms_predictions
 
-        class_name = prediction['class']
+# Fungsi untuk deteksi objek menggunakan Roboflow Model
+def detect_objects(image):
+    # Menyimpan gambar sementara
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
+        image.save(temp_file, format="JPEG")
+        temp_file_path = temp_file.name
 
-        # Choose a color for the class, if the class is not in the palette, use white
-        box_color = color_palette.get(class_name, (255, 255, 255))
+    # Slice gambar menjadi potongan-potongan kecil
+    slice_image_result = slice_image(
+        image=temp_file_path,
+        output_file_name="sliced_image",
+        output_dir="/tmp/sliced/",
+        slice_height=256,
+        slice_width=256,
+        overlap_height_ratio=0.1,
+        overlap_width_ratio=0.1
+    )
+
+    # Print to check the available attributes of the slice_image_result object
+    print(f"Slice result: {slice_image_result}")
+
+    # Try accessing the sliced image paths from the result object
+    try:
+        sliced_image_paths = slice_image_result.sliced_image_paths  # Assuming this is the correct attribute
+        print(f"Sliced image paths: {sliced_image_paths}")
+    except AttributeError:
+        print("Failed to access sliced_image_paths attribute.")
+        sliced_image_paths = []
+
+    # Menyimpan semua prediksi untuk setiap potongan gambar
+    all_predictions = []
+
+    # Prediksi pada setiap potongan gambar
+    for sliced_image_path in sliced_image_paths:
+        if isinstance(sliced_image_path, str):
+            predictions = model.predict(image_path=sliced_image_path).json()
+            all_predictions.extend(predictions['predictions'])
+        else:
+            print(f"Skipping invalid image path: {sliced_image_path}")
+    
+    # Aplikasikan NMS untuk menghapus duplikat deteksi
+    postprocessed_predictions = apply_nms(all_predictions, iou_threshold=0.5)
 
-        # Draw a bounding box around the detected object
-        cv2.rectangle(tiled_image, (x1, y1), (x2, y2), box_color, 2)  # Bounding box with thickness=2
+    # Annotate gambar dengan hasil prediksi
+    annotated_image = model.annotate_image_with_predictions(temp_file_path, postprocessed_predictions)
 
-        # Put the class name label on the bounding box
-        cv2.putText(tiled_image, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, box_color, 2)  # Label
+    # Simpan gambar hasil annotasi
+    output_image_path = "/tmp/prediction.jpg"
+    annotated_image.save(output_image_path)
 
-        # Count the class occurrences
+    # Menghitung jumlah objek per kelas
+    class_count = {}
+    for detection in postprocessed_predictions:
+        class_name = detection['class']
         if class_name in class_count:
             class_count[class_name] += 1
         else:
             class_count[class_name] = 1
 
-    # Create a result text to show class counts
-    result_text = "Object counts per class:\n"
+    # Hasil perhitungan objek
+    result_text = "Jumlah objek per kelas:\n"
     for class_name, count in class_count.items():
-        result_text += f"{class_name}: {count} objects\n"
+        result_text += f"{class_name}: {count} objek\n"
 
-    # Remove the temporary file after processing
+    # Hapus file sementara
     os.remove(temp_file_path)
 
-    return result_text  # Only return result_text for object counting
+    return output_image_path, result_text
 
-# Gradio Interface
+# Membuat antarmuka Gradio
 iface = gr.Interface(
-    fn=detect_objects,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(),  # Only output the text with object counts
-    live=True
+    fn=detect_objects,                         # Fungsi yang dipanggil saat gambar diupload
+    inputs=gr.Image(type="pil"),               # Input berupa gambar
+    outputs=[gr.Image(), gr.Textbox()],        # Output gambar dan teks
+    live=True                                    # Menampilkan hasil secara langsung
 )
 
-# Launch Gradio app
-iface.launch(debug=True)
+# Menjalankan antarmuka
+iface.launch()