import io import json import os import tempfile import zipfile import datetime import requests import random import gradio as gr from PIL import Image, ImageDraw, ImageFont API_URL = os.getenv("API_URL") API_KEY = os.getenv("API_KEY") MODEL = os.getenv("MODEL") ALLOWED_EXTENSIONS = ( ".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp", ".tif", ".tiff", ".avif", ".ico", ".ppm", ".pgm", ".pbm", ".eps", ".icns" ) def get_color_for_label(label, color_mapping): if label in color_mapping: return color_mapping[label] color = (random.randint(0,255), random.randint(0,255), random.randint(0,255)) color_mapping[label] = color return color def process_single_image(image, prompt, confidence, file_format): if "," in prompt: prompt_list = [p.strip() for p in prompt.split(",")] else: prompt_list = [prompt.strip()] prompt_list_lower = [p.lower() for p in prompt_list] buffered = io.BytesIO() image.save(buffered, format="PNG") buffered.seek(0) files = {"image": buffered} data = { "prompts": prompt_list, "model": MODEL, "confidence": confidence } headers = {"Authorization": "Basic " + API_KEY} response = requests.post(API_URL, files=files, data=data, headers=headers) response_data = response.json() print("Response data:", response_data) detections_list = response_data.get("data", []) if detections_list and isinstance(detections_list, list): detections = detections_list[0] else: detections = [] annotated_image = image.copy() draw = ImageDraw.Draw(annotated_image) try: font = ImageFont.truetype("arial.ttf", size=15) except IOError: font = ImageFont.load_default() color_mapping = {} for det in detections: box = det["bounding_box"] score = det["score"] detection_label = det.get("label", prompt).strip() box_color = get_color_for_label(detection_label, color_mapping) draw.rectangle(box, outline=box_color, width=2) text = f"{detection_label} {score:.2f}" bbox = draw.textbbox((0, 0), text, font=font) text_width = bbox[2] - bbox[0] text_height = bbox[3] - bbox[1] text_origin = (box[0], box[1] - text_height) draw.rectangle([text_origin, (box[0] + text_width, box[1])], fill=box_color) draw.text(text_origin, text, fill="white", font=font) img_width, img_height = image.size if file_format.lower() == "json": detection_info = {"detections": detections} ext = ".json" content = json.dumps(detection_info, indent=2) elif file_format.lower() == "txt": lines = [] for det in detections: detection_label = det.get("label", prompt).strip().lower() if detection_label not in prompt_list_lower: continue class_id = prompt_list_lower.index(detection_label) x1, y1, x2, y2 = det["bounding_box"] x_center = ((x1 + x2) / 2) / img_width y_center = ((y1 + y2) / 2) / img_height bbox_width = (x2 - x1) / img_width bbox_height = (y2 - y1) / img_height line = f"{class_id} {x_center:.6f} {y_center:.6f} {bbox_width:.6f} {bbox_height:.6f}" lines.append(line) content = "\n".join(lines) ext = ".txt" else: detection_info = {"detections": detections} ext = ".json" content = json.dumps(detection_info, indent=2) temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=ext, mode="w", encoding="utf-8") temp_file.write(content) temp_file.close() return annotated_image, temp_file.name def auto_annotate_batch(input_files, prompt, confidence, file_format): images_and_names = [] if len(input_files) == 1 and str(input_files[0]).lower().endswith(".zip"): zip_path = input_files[0] extract_dir = tempfile.mkdtemp() with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(extract_dir) for root, _, files in os.walk(extract_dir): for file in files: if file.lower().endswith(ALLOWED_EXTENSIONS): img_path = os.path.join(root, file) try: img = Image.open(img_path).convert("RGB") images_and_names.append((img, file)) except Exception as e: print(f"Gagal membuka {img_path}: {e}") else: for file_path in input_files: if file_path.lower().endswith(ALLOWED_EXTENSIONS): try: img = Image.open(file_path).convert("RGB") original_filename = os.path.basename(file_path) images_and_names.append((img, original_filename)) except Exception as e: print(f"Gagal membuka {file_path}: {e}") annotated_images = [] detection_file_paths = [] for img, original_filename in images_and_names: ann_img, temp_det_file = process_single_image(img, prompt, confidence, file_format) annotated_images.append(ann_img) base_name, _ = os.path.splitext(original_filename) new_name = base_name + (".json" if file_format.lower() == "json" else ".txt") new_path = os.path.join(os.path.dirname(temp_det_file), new_name) os.rename(temp_det_file, new_path) detection_file_paths.append(new_path) timestamp = datetime.datetime.now().strftime("%d-%m-%Y_%H-%M-%S") zip_filename = f"Data_Annotate_{timestamp}.zip" zip_out_path = os.path.join(tempfile.gettempdir(), zip_filename) with zipfile.ZipFile(zip_out_path, 'w') as zipf: for file_path in detection_file_paths: zipf.write(file_path, os.path.basename(file_path)) return annotated_images, zip_out_path, detection_file_paths iface = gr.Interface( fn=auto_annotate_batch, inputs=[ gr.File(file_count="multiple", type="filepath", label="Upload Image (multiple files or ZIP)"), gr.Textbox(lines=1, placeholder="Enter object prompt (separate with comma if more than one)", label="Prompt"), gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.25, label="Confidence Threshold"), gr.Radio(choices=["json", "txt"], value="json", label="Format Annotation File") ], outputs=[ gr.Gallery(label="Annotated Images"), gr.File(label="Download All Annotations (ZIP)"), gr.File(label="Download Individual Annotation Files") ], title="Auto Annotate") if __name__ == "__main__": iface.launch(debug=True)