Update app.py

app.py

@@ -1,121 +1,103 @@
 import gradio as gr
-from roboflow import Roboflow
-import tempfile
 import os
-from sahi.slicing import slice_image
-import numpy as np
+import tempfile
+import math
 import cv2
+import numpy as np
+import supervision as sv
+from roboflow import Roboflow
 
-# Inisialisasi Roboflow (for model path)
+# Initialize Roboflow
 rf = Roboflow(api_key="Otg64Ra6wNOgDyjuhMYU")
 project = rf.workspace("alat-pelindung-diri").project("nescafe-4base")
 model = project.version(16).model
 
-# Fungsi untuk melakukan Non-Maximum Suppression (NMS)
-def apply_nms(predictions, iou_threshold=0.5):
-    boxes = []
-    scores = []
-    classes = []
+# 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
 
-    # Extract boxes, scores, and class info
-    for prediction in predictions:
-        boxes.append(prediction['bbox'])
-        scores.append(prediction['confidence'])
-        classes.append(prediction['class'])
-
-    boxes = np.array(boxes)
-    scores = np.array(scores)
-    classes = np.array(classes)
+# 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
 
-    # Perform NMS using OpenCV
-    indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), score_threshold=0.25, nms_threshold=iou_threshold)
-    nms_predictions = []
+    image_shape = (img.shape[1], img.shape[0])
 
-    for i in indices.flatten():
-        nms_predictions.append({
-            'class': classes[i],
-            'bbox': boxes[i],
-            'confidence': scores[i]
-        })
+    # 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)
 
-    return nms_predictions
+    # 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()
 
-# Fungsi untuk deteksi objek menggunakan Roboflow Model
-def detect_objects(image):
-    # Menyimpan gambar sementara
+    # 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
 
-    # 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)
+    # 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 = {}
 
-    # Annotate gambar dengan hasil prediksi
-    annotated_image = model.annotate_image_with_predictions(temp_file_path, postprocessed_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
+    }
 
-    # Simpan gambar hasil annotasi
-    output_image_path = "/tmp/prediction.jpg"
-    annotated_image.save(output_image_path)
+    # 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)
 
-    # Menghitung jumlah objek per kelas
-    class_count = {}
-    for detection in postprocessed_predictions:
-        class_name = detection['class']
+        class_name = prediction['class']
+
+        # 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))
+
+        # Draw a bounding box around the detected object
+        cv2.rectangle(tiled_image, (x1, y1), (x2, y2), box_color, 2)  # Bounding box with thickness=2
+
+        # 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
+
+        # Count the class occurrences
         if class_name in class_count:
             class_count[class_name] += 1
         else:
             class_count[class_name] = 1
 
-    # Hasil perhitungan objek
-    result_text = "Jumlah objek per kelas:\n"
+    # Create a result text to show class counts
+    result_text = "Object counts per class:\n"
     for class_name, count in class_count.items():
-        result_text += f"{class_name}: {count} objek\n"
+        result_text += f"{class_name}: {count} objects\n"
 
-    # Hapus file sementara
+    # Remove the temporary file after processing
     os.remove(temp_file_path)
 
-    return output_image_path, result_text
+    return result_text  # Only return result_text for object counting
 
-# Membuat antarmuka Gradio
+# Gradio Interface
 iface = gr.Interface(
-    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
+    fn=detect_objects,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Textbox(),  # Only output the text with object counts
+    live=True
 )
 
-# Menjalankan antarmuka
-iface.launch()
+# Launch Gradio app
+iface.launch(debug=True)