File size: 5,363 Bytes
75b9799 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
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
# ========== Konfigurasi ==========
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"))
# OWLv2 API Config
OWLV2_API_URL = "https://api.landing.ai/v1/tools/text-to-object-detection"
OWLV2_PROMPTS = ["beverage", "bottle", "cans", "boxed milk", "milk"]
# Inisialisasi Model YOLO
rf = Roboflow(api_key=rf_api_key)
project = rf.workspace(workspace).project(project_name)
yolo_model = project.version(model_version).model
# ========== Fungsi Deteksi Kombinasi ==========
def detect_combined(image):
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
image.save(temp_file, format="JPEG")
temp_path = temp_file.name
try:
# ========== [1] YOLO: Deteksi Produk Nestlé (Per Class) ==========
yolo_pred = yolo_model.predict(temp_path, confidence=50, overlap=80).json()
nestle_class_count = {}
nestle_boxes = []
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']))
total_nestle = sum(nestle_class_count.values())
# ========== [2] OWLv2: Deteksi Kompetitor ==========
with open(temp_path, "rb") as image_file:
response = requests.post(OWLV2_API_URL,
files={"image": image_file},
data={"prompts": OWLV2_PROMPTS, "model": "owlv2"})
owlv2_pred = response.json().get("objects", [])
competitor_class_count = {}
competitor_boxes = []
for obj in owlv2_pred:
x1, y1, x2, y2 = obj["bbox"]
class_name = obj["label"].strip().lower()
if not is_overlap((x1, y1, x2, y2), nestle_boxes):
competitor_class_count[class_name] = competitor_class_count.get(class_name, 0) + 1
competitor_boxes.append({"class": class_name, "box": (x1, y1, x2, y2), "confidence": obj["score"]})
total_competitor = sum(competitor_class_count.values())
# ========== [3] Format Output ==========
result_text = "Product Nestle\n\n"
for class_name, count in nestle_class_count.items():
result_text += f"{class_name}: {count}\n"
result_text += f"\nTotal Products Nestle: {total_nestle}\n\n"
result_text += f"Total Unclassified Products: {total_competitor}\n" if competitor_class_count else "No Unclassified Products detected\n"
# ========== [4] Visualisasi ==========
img = cv2.imread(temp_path)
for pred in yolo_pred['predictions']:
x, y, w, h = pred['x'], pred['y'], pred['width'], pred['height']
cv2.rectangle(img, (int(x-w/2), int(y-h/2)), (int(x+w/2), int(y+h/2)), (0,255,0), 2)
cv2.putText(img, pred['class'], (int(x-w/2), int(y-h/2-10)), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0,255,0), 2)
for comp in competitor_boxes:
x1, y1, x2, y2 = comp['box']
display_name = "unclassified" if comp['class'] in OWLV2_PROMPTS else comp['class']
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
cv2.putText(img, f"{display_name} {comp['confidence']:.2f}", (int(x1), int(y1-10)), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 255), 2)
output_path = "/tmp/combined_output.jpg"
cv2.imwrite(output_path, img)
return output_path, result_text
except Exception as e:
return temp_path, f"Error: {str(e)}"
finally:
os.remove(temp_path)
def is_overlap(box1, boxes2, threshold=0.3):
x1_min, y1_min, x1_max, y1_max = box1
for b2 in boxes2:
x2, y2, w2, h2 = b2
x2_min = x2 - w2/2
x2_max = x2 + w2/2
y2_min = y2 - h2/2
y2_max = y2 + 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_overlap / area_box1 > threshold:
return True
return False
with gr.Blocks(theme=gr.themes.Base(primary_hue="teal", secondary_hue="teal", neutral_hue="slate")) as iface:
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")
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]
) |