import gradio as gr from dotenv import load_dotenv from roboflow import Roboflow import tempfile import os import requests import cv2 import numpy as np import subprocess # ========== Load Environment Variables ========== load_dotenv() # Roboflow Config rf_api_key = os.getenv("ROBOFLOW_API_KEY") workspace = os.getenv("ROBOFLOW_WORKSPACE") project_name = os.getenv("ROBOFLOW_PROJECT") model_version = int(os.getenv("ROBOFLOW_MODEL_VERSION")) # CountGD Config COUNTGD_API_KEY = os.getenv("COUNTGD_API_KEY") # Inisialisasi YOLO Model dari Roboflow rf = Roboflow(api_key=rf_api_key) project = rf.workspace(workspace).project(project_name) yolo_model = project.version(model_version).model # ========== Fungsi untuk Mengecek Overlap antara YOLO dan CountGD ========== def is_overlap(box1, boxes2, threshold=0.5): """ Mengecek apakah box1 (format: (x_min, y_min, x_max, y_max)) overlap dengan salah satu box di boxes2 (format: (x_center, y_center, width, height)) berdasarkan IoU, menggunakan threshold yang ditetapkan. """ x1_min, y1_min, x1_max, y1_max = box1 for b2 in boxes2: x_center, y_center, w2, h2 = b2 x2_min = x_center - w2 / 2 x2_max = x_center + w2 / 2 y2_min = y_center - h2 / 2 y2_max = y_center + h2 / 2 dx = min(x1_max, x2_max) - max(x1_min, x2_min) dy = min(y1_max, y2_max) - max(y1_min, y2_min) if dx > 0 and dy > 0: area_overlap = dx * dy area_box1 = (x1_max - x1_min) * (y1_max - y1_min) if area_box1 > 0 and (area_overlap / area_box1) > threshold: return True return False # ========== Fungsi untuk Menghitung IoU antar dua bounding box ========== def iou(boxA, boxB): """ Menghitung Intersection over Union (IoU) antara dua bounding box. Masing-masing box dalam format (x_min, y_min, x_max, y_max). """ xA = max(boxA[0], boxB[0]) yA = max(boxA[1], boxB[1]) xB = min(boxA[2], boxB[2]) yB = min(boxA[3], boxB[3]) interArea = max(0, xB - xA) * max(0, yB - yA) boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1]) boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1]) return interArea / float(boxAArea + boxBArea - interArea) if (boxAArea + boxBArea - interArea) > 0 else 0 # ========== Fungsi Deteksi Kombinasi ========== def detect_combined(image): # Simpan image ke file sementara with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file: image.save(temp_file, format="JPEG") temp_path = temp_file.name try: # ===== YOLO Detection ===== yolo_pred = yolo_model.predict(temp_path, confidence=50, overlap=80).json() # 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" headers = {"Authorization": f"Basic {COUNTGD_API_KEY}"} competitor_boxes = [] COUNTGD_PROMPTS = ["cans", "bottle", "boxed milk", "milk"] for prompt in COUNTGD_PROMPTS: with open(temp_path, "rb") as f: files = {"image": f} data = {"prompts": [prompt], "model": "owlv2"} response = requests.post(url, files=files, data=data, headers=headers) result = response.json() if 'data' in result and result['data']: detections = result['data'][0] for obj in detections: if 'bounding_box' in obj: x1, y1, x2, y2 = obj['bounding_box'] countgd_box = (x1, y1, x2, y2) # 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 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: duplicate = True break if not duplicate: competitor_boxes.append(countgd_box) # ===== Visualisasi ===== img = cv2.imread(temp_path) # Gambar bounding box YOLO (hijau) for pred in yolo_pred['predictions']: x, y, w, h = pred['x'], pred['y'], pred['width'], pred['height'] pt1 = (int(x - w/2), int(y - h/2)) pt2 = (int(x + w/2), int(y + h/2)) cv2.rectangle(img, pt1, pt2, (0, 255, 0), 2) cv2.putText(img, pred['class'], (pt1[0], pt1[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,255,0), 3) # Gambar bounding box CountGD (merah) for box in competitor_boxes: x1, y1, x2, y2 = box cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2) cv2.putText(img, "unclassified", (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,0,255), 3) output_path = "/tmp/combined_output.jpg" cv2.imwrite(output_path, img) # Buat result text untuk count produk Nestlé per class dan total keseluruhan result_text = "Produk Nestlé:\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: return temp_path, f"Error: {str(e)}" finally: if os.path.exists(temp_path): os.remove(temp_path) # ========== Gradio Interface ========== with gr.Blocks(theme=gr.themes.Base(primary_hue="teal", secondary_hue="teal", neutral_hue="slate")) as iface: gr.Markdown("""

NESTLE - STOCK COUNTING

""") with gr.Row(): with gr.Column(): input_image = gr.Image(type="pil", label="Input Image") with gr.Column(): output_image = gr.Image(label="Detect Object") with gr.Column(): output_text = gr.Textbox(label="Counting Object") # Tombol untuk memproses input detect_button = gr.Button("Detect") # Hubungkan tombol dengan fungsi deteksi detect_button.click( fn=detect_combined, inputs=input_image, outputs=[output_image, output_text] ) iface.launch()