Update app.py

app.py

@@ -74,8 +74,13 @@ def detect_combined(image):
     try:
         # ===== YOLO Detection =====
         yolo_pred = yolo_model.predict(temp_path, confidence=50, overlap=80).json()
-        # Get YOLO boxes as (x_center, y_center, width, height)
-        nestle_boxes = [(pred['x'], pred['y'], pred['width'], pred['height']) for pred in yolo_pred['predictions']]
+        # Hitung bounding box dan count per class untuk produk Nestlé
+        nestle_boxes = []
+        nestle_class_count = {}
+        for pred in yolo_pred['predictions']:
+            class_name = pred['class']
+            nestle_class_count[class_name] = nestle_class_count.get(class_name, 0) + 1
+            nestle_boxes.append((pred['x'], pred['y'], pred['width'], pred['height']))
         
         # ===== CountGD Detection =====
         url = "https://api.landing.ai/v1/tools/text-to-object-detection"
@@ -99,7 +104,7 @@ def detect_combined(image):
                         # Prioritaskan deteksi YOLO: hapus jika overlap dengan YOLO (threshold 0.5)
                         if is_overlap(countgd_box, nestle_boxes, threshold=0.5):
                             continue
-                        # Hindari duplikasi antar CountGD: jika IoU dengan deteksi lain > 0.4, lewati
+                        # Hindari duplikasi antar deteksi CountGD: jika IoU dengan deteksi lain > 0.4, lewati
                         duplicate = False
                         for existing_box in competitor_boxes:
                             if iou(countgd_box, existing_box) > 0.4:
@@ -128,8 +133,14 @@ def detect_combined(image):
         output_path = "/tmp/combined_output.jpg"
         cv2.imwrite(output_path, img)
         
-        # Buat result text untuk counting object
-        result_text = f"Total Produk Nestlé: {len(nestle_boxes)}\nTotal Unclassified Products: {len(competitor_boxes)}"
+        # Buat result text untuk count produk Nestlé per class dan total keseluruhan
+        result_text = "Produk Nestlé per Class:\n"
+        for class_name, count in nestle_class_count.items():
+            result_text += f"{class_name}: {count}\n"
+        total_nestle = sum(nestle_class_count.values())
+        result_text += f"\nTotal Produk Nestlé: {total_nestle}\n"
+        result_text += f"Total Unclassified Products: {len(competitor_boxes)}"
+        
         return output_path, result_text
     
     except Exception as e:
@@ -144,17 +155,13 @@ with gr.Blocks(theme=gr.themes.Base(primary_hue="teal", secondary_hue="teal", ne
     gr.Markdown("""<div style="text-align: center;"><h1>NESTLE - STOCK COUNTING</h1></div>""")
     
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Input Image")
+        input_image = gr.Image(type="pil", label="Input Image")
     with gr.Row():
-        with gr.Column():
-            detect_image_button = gr.Button("Detect Image")
+        detect_image_button = gr.Button("Detect Image")
     with gr.Row():
-        with gr.Column():
-            output_image = gr.Image(label="Detect Object")
+        output_image = gr.Image(label="Detect Object")
     with gr.Row():
-        with gr.Column():
-            output_text = gr.Textbox(label="Counting Object")
+        output_text = gr.Textbox(label="Counting Object")
     
     detect_image_button.click(fn=detect_combined, inputs=input_image, outputs=[output_image, output_text])